FastChain против графических процессоров в DiffEqFlux

Question

Для обучения модели на GPU я использую

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

против

Обучение ЦП

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

Более 1000 итераций Fastchain на порядки быстрее, чем GPU Tesla K40c. Это ожидаемое поведение? В противном случае, мог ли я что-то не так с реализацией модели на графических процессорах. MWE для реализации на GPU следующим образом:

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

Эта модель слишком мала для параллелизма графического процессора, чтобы действительно иметь значение. Нейронная сеть состоит из трех матвеков, 100x3, 100x100, 3x100. Единственное ядро ​​с ядром, которое, вероятно, приближается к безубыточности, - это средний, где матрица 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)

Таким образом, хотя ускорение самой большой операции действительно существует, этого недостаточно, чтобы преодолеть замедление, перенеся другие небольшие операции на графический процессор. Это связано с тем, что графические процессоры имеют высокий этаж (на моей машине около 12 с), поэтому вы должны убедиться, что ваша проблема достаточно велика, чтобы она действительно имела смысл. Как правило, машинное обучение выигрывает от графических процессоров, потому что в нем преобладают большие матричные умножения со слоями размером в десятки тысяч.