记录一次Julia手撕高斯过程分类的经验

高斯过程分类

本文章不太涉及理论上的问题,而是讨论在已有算法的情况下,如何从Julia工程实践的角度加速我们的数值计算,可以说Julia在加速前后的差距已经类似Python和C++的速度差距了。

由于Julia不是面向对象语言,因此在写有关机器学习的库时,很难效仿 scikit-learn 的接口,在Julia当中只能使用类似的结构体:

Julia
mutable struct GPC{K<:Kernel}
    kernel::K
    X::Matrix{Float64}          # 训练输入
    y::Vector{Float64}          # 训练标签 ±1
    α::Vector{Float64}          # 预测时用的系数
    L::Matrix{Float64}          # Cholesky 分解 (I + W^{1/2} K W^{1/2})
    f_hat::Vector{Float64}      # 后验众数
    W::Vector{Float64}          # 负 Hessian 对角线
end

# 下面是构造函数
function GPC(kernel::K) where {K<:Kernel}
    return GPC{K}(
        kernel,
        Matrix{Float64}(undef,0,0),
        Float64[],
        Float64[],
        Matrix{Float64}(undef,0,0),
        Float64[],
        Float64[]
    )
end

对于Kernel,我们需要预先定义,考虑使用RBF核,此时为适应核矩阵的计算,需要通过多重派发实现:

Julia
begin
    abstract type Kernel end

    struct SEKernel <: Kernel
        σ::Float64   # 信号方差
::Float64   # 长度尺度
    end

    # 有默认值的析构函数
    SEKernel(; σ=1.0, ℓ=1.0) = SEKernel(σ, ℓ)


    # 核函数:单个样本对单个样本
    function::SEKernel)(x1::V1, x2::V2) where {V1<:AbstractVector, V2<:AbstractVector}
        sqdist = norm(x1 - x2, 2)
        return κ.σ^2 * exp(-sqdist / (2 * κ.^2))
    end

    # 核函数:单个样本对一批样本
    function::SEKernel)(x1::V, x2::M) where {V<:AbstractVector, M<:AbstractMatrix}
        m = size(x2, 1)

        K = Vector{Float64}(undef, m)

        for j in 1:m
            row_j = @view x2[j, :]
            K[j] = κ(x1, row_j)
        end

        return K
    end
    
    # 核函数:一批样本对一批样本
    function::SEKernel)(x1::M1, x2::M2) where {M1<:AbstractMatrix, M2<:AbstractMatrix}
        n = size(x1, 1)
        m = size(x2, 1)

        K = Matrix{Float64}(undef, n, m)
        for i in 1:n
            row_i = @view x1[i, :]
            for j in 1:m
                row_j = @view x2[j, :]
                K[i,j] = κ(row_i, row_j)
            end
        end
        return K
    end
end

上述就是基本的接口定义,此时为了解决任务,还需要如下的方法:

  • fit! :就地拟合模型的接口,高斯过程二分类的实现相对比较复杂,可能需要理解如下的一些概念:
    • Laplace 近似
    • 期望传播
    • 变分推断方法
    • Cholesky 分解
  • predict :预测数据的接口,相对简单,就是需要利用数值积分,在本文当中为简便起见使用 QuadGK.jl 库来实现简单的数值积分.

上述的写法有一些利用了 Julia 特性的结构,其中一个就是类型注解,注意:

Julia
function::SEKernel)(x1::V1, x2::V2) where {V1<:AbstractVector, V2<:AbstractVector}

Julia
function::SEKernel)(x1::AbstractVector, x2::AbstractVector)

二者之间具有本质上的差距,因为 Julia 会根据具体的数值类型来特定优化,前者要更好,后者在计算时编译器无法识别函数当中的 x1 是什么类型,可能无法调用 BLAS 等常用高性能计算工具。

拟合函数

拟合函数 fit! 的计算是高斯过程分类实现的核心,也是计算耗时的大头,下面我们直接给出代码逐个解析:

Julia
function fit!(model::GPC,
			  X::M,
			  y::T,
			  max_iter::Int64 = 1000,
			  alpha::Float64 = 0.05,
			  tol::Float64 = 1e-5;
			  verbose::Bool = false) where {T<:AbstractVector,M<:AbstractMatrix}

	X_train = Matrix{Float64}(X)
	y_train = Vector{Float64}(y)

	K = model.kernel(X_train, X_train)
	@inbounds for i in axes(K, 1)
		K[i, i] += 1e-6
	end

	n = length(y_train)
	f = zeros(Float64, n)
	grad = zeros(Float64, n)
	w = zeros(Float64, n)
	sqrt_w = zeros(Float64, n)
	a = zeros(Float64, n)
	rhs = zeros(Float64, n)
	direction = zeros(Float64, n)
	B = similar(K)

	err = tol + 1
	iter = 0
	L = Matrix{Float64}(undef, 0, 0)

	while (err > tol) && (iter < max_iter)
		if verbose
			println("第 $(iter + 1) 次迭代")
		end

		neg_hessian_diag!(w, f)
		@. sqrt_w = sqrt(w)

		# 数值稳定化并求解牛顿方向
		copyto!(B, K)
		rmul!(B, Diagonal(sqrt_w))
		lmul!(Diagonal(sqrt_w), B)

		@inbounds for i in axes(B, 1)
			B[i, i] += 1.0 + 1e-8
		end

		F = cholesky!(Hermitian(B))
		L = Matrix(F.L)

		grad_log_likelihood!(grad, y_train, f)
		@. rhs = w * f + grad
		mul!(a, K, rhs)
		@. a = sqrt_w * a
		ldiv!(LowerTriangular(L), a)
		ldiv!(UpperTriangular(L'), a)
		@. a = sqrt_w * a

		mul!(direction, K, a)

		# 更新迭代状态,避免生成新的 f 向量
		err_sq = 0.0
		@inbounds for i in eachindex(f, direction)
			step = direction[i] - f[i]
			new_f = f[i] + alpha * step
			diff = f[i] - new_f
			err_sq += diff * diff
			f[i] = new_f
		end

		err = sqrt(err_sq)
		iter += 1

		if verbose
			println("err = $err")
		end
	end

	# 收敛后将结果写回模型
	grad_log_likelihood!(grad, y_train, f)
	model.f_hat = copy(f)
	model.W = copy(w)
	model.α = copy(grad)   # ∇ log p(y|f_hat)
	model.L = L
	model.X = X_train
	model.y = y_train

	return model

end

对类型的特化

Julia
function fit!(model::GPC,
			  X::M,
			  y::T,
			  max_iter::Int64 = 1000,
			  alpha::Float64 = 0.05,
			  tol::Float64 = 1e-5;
			  verbose::Bool = false) where {T<:AbstractVector,M<:AbstractMatrix}

	X_train = Matrix{Float64}(X)
	y_train = Vector{Float64}(y)

开头我们将 X,y 进行了到矩阵和向量的强制转化,这样之后的计算当中,所有的计算都将调用合适的线性代数库,速度也将大大提升。

关闭边界检查

在我们的算法里面,很多时候我们的循环不会触发越界,因此可以关闭边界检查来实现加速:

Julia
@inbounds for i in axes(K, 1)
    K[i, i] += 1e-6
end

这里是对矩阵 K 加上了稳定的对角线因子,相对简单。

矩阵预分配

预分配是一个重要环节,对于循环里面会反复用到的中间量,一般的建议是提前分配好空间,这样可以降低 GC 的压力,否则循环当中会反复创建新的向量和矩阵,对于稠密矩阵计算这将会耗费巨量的内存。

Julia
n = length(y_train)
f = zeros(Float64, n)
grad = zeros(Float64, n)
w = zeros(Float64, n)
sqrt_w = zeros(Float64, n)
a = zeros(Float64, n)
rhs = zeros(Float64, n)
direction = zeros(Float64, n)
B = similar(K)

这种预分配虽然好,但是后续的计算也必须要搭配上,因为常规的矩阵乘法和 Cholesky 分解并不是就地操作,需要换为不会产生新矩阵的就地操作,因此后面的代码当中:

  • mul!(a, K, rhs) 替代 a = K * rhs
  • ldiv!(LowerTriangular(L), a) 替代 L \ a ,这里 LowerTriangular 实际上还指示了是下三角方程组,因此计算起来很快速;
  • cholesky!(Hermitian(B)) 代替 cholesky(B)

上述文档详见 Julia LinearAlgebra 官方文档 ,这些优化可以算是数值计算的基本功了。

Leave a Comment

您的邮箱地址不会被公开。 必填项已用 * 标注