FastChain против графических процессоров в DiffEqFlux - PullRequest
Новая
       54

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

1 голос
/ 16 апреля 2020

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

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

против

Обучение ЦП 

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

За 1000 итераций Fastchain на несколько порядков быстрее, чем работает графический процессор Tesla K40 c. Это ожидаемое поведение? В противном случае, я мог бы сделать что-то не так с реализацией модели на графических процессорах. 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)

1 Ответ

2 голосов
/ 16 апреля 2020

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

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

...