DiffEqFlux 中的 FastChain 与 GPU

Question

对于 model 的 GPU 训练，我正在使用

dudt = Chain(Dense(3,100,tanh),
    Dense(100,3)) |> gpu

相对

CPU 训练

dudt = FastChain(   
              FastDense(3,100,tanh),
              FastDense(100,3))

超过 1000 次迭代，Fastchain 比运行 GPU Tesla K40c 快几个数量级。 这是预期的行为吗？ 否则，我会在 GPU 上实现 model 做错什么吗？ GPU 的 MWE 实现如下：

function lorenz(du,u,p,t)
    σ = p[1]; ρ = p[2]; β = p[3]
    du[1] = σ*(u[2]-u[1])
    du[2] = u[1]*(ρ-u[3]) - u[2]
    du[3] = u[1]*u[2] - β*u[3]
    return 
end
u0 = Float32[1.0,0.0,0.0]               
tspan = (0.0,1.0)                      
para = [10.0,28.0,8/3]                      
prob = ODEProblem(lorenz, u0, tspan, para)  
t = range(tspan[1],tspan[2],length=101)
ode_data = Array(solve(prob,Tsit5(),saveat=t))
ode_data = cu(ode_data)

u0train = [1.0,0.0,0.0] |> gpu
tspantrain = (0.0,1.0)  
ttrain = range(tspantrain[1],tspantrain[2],length=101)  
dudt = Chain(Dense(3,100,tanh),
    Dense(100,3)) |> gpu
n_ode = NeuralODE((dudt),tspantrain,Tsit5(),saveat=ttrain)

function predict_n_ode(p)
  n_ode(u0train,p)
end

function loss_n_ode(p)
    pred = predict_n_ode(p) |> gpu
    loss = sum(abs2, pred .- ode_data)
    loss,pred
end

res1 = DiffEqFlux.sciml_train(loss_n_ode, n_ode.p, ADAM(0.01), cb=cb, maxiters = 1000)

Answer 1

model 对于 GPU 并行性来说太小而无法真正发挥作用。 神经网络本质上是一个 3 个 matvecs，100x3、100x100、3x100。 唯一一个可能接近收支平衡的 kernel 是中间那个，其中一个 100x100 矩阵乘以一个长度为 100 的向量。

例如，在我的机器上：

using BenchmarkTools, CuArrays
A = rand(100,100); x = rand(100);
@btime A*x; # 56.299 μs (1 allocation: 896 bytes)
gA = cu(A); gx = cu(x)
@btime gA*gx; # 12.499 μs (6 allocations: 160 bytes)

A = rand(100,3); x = rand(3);
@btime A*x; # 251.695 ns (1 allocation: 896 bytes)
gA = cu(A); gx = cu(x)
@btime gA*gx; # 12.212 μs (6 allocations: 160 bytes)

因此，虽然最大操作的加速确实存在，但通过在 GPU 上进行其他小操作来克服减速是不够的。 这是因为 GPU 有很高的底限（在我的机器上大约 12μs），所以你必须确保你的问题足够大，才能真正有意义。 通常，机器学习受益于 GPU，因为它主要由具有数万层大小的大型矩阵乘法控制。