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

高斯过程分类
本文章不太涉及理论上的问题,而是讨论在已有算法的情况下,如何从Julia工程实践的角度加速我们的数值计算,可以说Julia在加速前后的差距已经类似Python和C++的速度差距了。
由于Julia不是面向对象语言,因此在写有关机器学习的库时,很难效仿 scikit-learn 的接口,在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核,此时为适应核矩阵的计算,需要通过多重派发实现:
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 特性的结构,其中一个就是类型注解,注意:
function (κ::SEKernel)(x1::V1, x2::V2) where {V1<:AbstractVector, V2<:AbstractVector}和
function (κ::SEKernel)(x1::AbstractVector, x2::AbstractVector)二者之间具有本质上的差距,因为 Julia 会根据具体的数值类型来特定优化,前者要更好,后者在计算时编译器无法识别函数当中的 x1 是什么类型,可能无法调用 BLAS 等常用高性能计算工具。
拟合函数
拟合函数 fit! 的计算是高斯过程分类实现的核心,也是计算耗时的大头,下面我们直接给出代码逐个解析:
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对类型的特化
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 进行了到矩阵和向量的强制转化,这样之后的计算当中,所有的计算都将调用合适的线性代数库,速度也将大大提升。
关闭边界检查
在我们的算法里面,很多时候我们的循环不会触发越界,因此可以关闭边界检查来实现加速:
@inbounds for i in axes(K, 1)
K[i, i] += 1e-6
end这里是对矩阵 K 加上了稳定的对角线因子,相对简单。
矩阵预分配
预分配是一个重要环节,对于循环里面会反复用到的中间量,一般的建议是提前分配好空间,这样可以降低 GC 的压力,否则循环当中会反复创建新的向量和矩阵,对于稠密矩阵计算这将会耗费巨量的内存。
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 官方文档 ,这些优化可以算是数值计算的基本功了。