Сравнение обучения нейросетей на CPU и GPU

Вернемся к примеру обучения многослойной нейросети задаче регрессии и сравним, как ускорится процесс при переходе на GPU.

Установка библиотек

Следующая команда может занять несколько минут, если названные библиотеки еще не были установлены.

Pkg.add(["CUDA", "Flux"])

Pkg.add( ["Flux", "CUDA", "cuDNN", "Random"] )

using Flux, CUDA

Подготовка

Создадим довольно большие векторы обучающих данных типа Float32:

Nx1, Nx2 = 200, 200
x1 = Float32.( range( -3, 3, length=Nx1 ) )
x2 = Float32.( range( -3, 3, length=Nx2 ) )
Xs = [ repeat( x1, outer=Nx2)  repeat( x2, inner=Nx1) ];
Ys = @. 3*(1-Xs[:,1])^2*exp(-(Xs[:,1]^2) - (Xs[:,2]+1)^2) - 10*(Xs[:,1]/5 - Xs[:,1]^3 - Xs[:,2]^5)*exp(-Xs[:,1]^2-Xs[:,2]^2) - 1/3*exp(-(Xs[:,1]+1) ^ 2 - Xs[:,2]^2);

И зададим достаточно большое время обучения:

epochs = 2000;
learning_rate = 0.08;

При напуске на CPU стоит начать с epochs=100 и медленно наращивать эту величину. Даже если мы далеки от "промышленных" значений, мы всё равно рискуем потерять очень много времени на CPU пока будем обучать нейросеть в неоптимальной постановке.

Загрузка данных и обучение на CPU

В качестве доработки предшествующего примера мы создаем процедуру обучения на основе DataLoader. С его помощью можно обучать нейросеть на батчах данных (подвыборках), а не подавать в вычислитель один пример за другим. Правильная настройка batchsize определяет производительность.

model_cpu = Chain( Dense( 2 => 20, relu ), Dense( 20 => 5, relu ), Dense( 5 => 1 ) )
model_copy = deepcopy( model_cpu ) # На будущее нам будет нужна точно такая же модель

data = [ (Xs', Ys') ]
loss( x, y ) = Flux.mse( model_cpu( x ), y )
opt = Adam( learning_rate );
ps = Flux.params( model_cpu );

using Random
Random.seed!( 1 ); # Установим генератор случайных чисел в воспроизводимое состояние

@time (for i in 1:epochs
    Flux.train!( loss, ps, data, opt )
end)

490.725439 seconds (17.02 M allocations: 42.837 GiB, 1.44% gc time, 2.05% compilation time)

Обучение длилось довольно долго. Если бы мы были ограничены CPU в своей работе, нам бы пришлось снизить количество точек в обучающей выборке, по крайней мере на первых порах обучения, пока мы подбирали бы гиперпараметры нейросети.

1-2% этого времени уходит на очистку памяти (gc_time = garbage collection) и на компиляцию кода (compile time).

Перенос процесса на GPU

При подключении библиотеки Flux, в нашем распоряжении появляется конструкция |> gpu (передать левое выражение в функцию gpu()). Она позволяет наглядно, без создания дополнительных уровней вложенности обработки данных, отправить матрицу или структуру на GPU и выполнить код на ней.

if CUDA.functional()
    model_gpu = model_copy |> gpu;
    
    Xg = Xs' |> gpu;
    Yg = Ys' |> gpu;
    
    data = [ (Xg, Yg) ]
    
    loss( x, y ) = Flux.mse( model_gpu( x ), y );
    opt = Adam( learning_rate );
    ps = Flux.params( model_gpu );
end;

Отметим, что мы пробовали использовать DataLoader. Его применение намного замедляло выполнение на GPU, из чего был сделан вывод о том, что стоит доверять процедуре распараллеливания, которую реализует Julia GPU Compiler.

Random.seed!( 1 ); # Установим генератор случайных чисел в воспроизводимое состояние

if CUDA.functional()
    @time CUDA.@sync ( for i in 1:epochs Flux.train!( loss, ps, data, opt ); end )
end

  3.122127 seconds (7.47 M allocations: 210.440 MiB, 16.94% gc time, 2.13% compilation time: 100% of which was recompilation)

Резюме по переносу на GPU

В целом, перенос нейросети на GPU требует трех дополнительных действий:

• переноса матрицы признаков Xs |> gpu
• переноса матрицы откликов Ys |> gpu
• переноса структуры с нейросетью model |> gpu

При небольшом количестве обучающих данных (100-200 примеров) нейросеть быстрее учится на CPU. Но на 40000 примерах преимущество на стороне GPU.

С поправкой на компиляцию (compile time) и очистку памяти после выполнения (gc time = garbage collection time) можно сказать, что

видеокарта справлялась с этой задачей примерно в 200 раз быстрее чем центральные процессоры.

Проверка результатов

Используем нейросеть, обученную на GPU, чтобы выполнить интерполяцию исходного набора данных:

gr()

# Вернем модель на CPU
if CUDA.functional()
    model_gpu = model_gpu |> cpu;
else
    model_gpu = model_cpu
end

# Создадим набор данных поменьше (иначе отрисовка займет очень много времени)
Nx1, Nx2 = 40, 40
x1 = Float32.( range( -3, 3, length=Nx1 ) )
x2 = Float32.( range( -3, 3, length=Nx2 ) )
Xs = [ repeat( x1, outer=Nx2)  repeat( x2, inner=Nx1) ];
Ys = @. 3*(1-Xs[:,1])^2*exp(-(Xs[:,1]^2) - (Xs[:,2]+1)^2) - 10*(Xs[:,1]/5 - Xs[:,1]^3 - Xs[:,2]^5)*exp(-Xs[:,1]^2-Xs[:,2]^2) - 1/3*exp(-(Xs[:,1]+1) ^ 2 - Xs[:,2]^2);

plot(
    surface( Xs[:,1], Xs[:,2], vec(Ys), c=:viridis, cbar=:false, title="Обучающая выборка", titlefont=font(10)),
    wireframe( x1, x2, vec(model_cpu( Xs' )), title="Нейросеть (CPU)", titlefont=font(10) ), 
    wireframe( x1, x2, vec(model_gpu( Xs' )), title="Нейросеть (GPU)", titlefont=font(10) ), 
    layout=(1,3), size=(1000,400)
)

Интересно, что одной установки Random.seed хватило, чтобы получить два идентичных (на вид) результата – один на CPU, другой на GPU.

Заключение

Мы обучили нейросеть интерполировать выборку в 40 000 элементов на CPU и на GPU.

Процесс обучения в течение 2 000 эпох на GPU был осуществлен почти в 200 раз быстрее, чем на CPU, что позволяет выполнить гораздо больше итераций по гиперпараметрам и обучать сети, гораздо больше по размеру, экономя время проектировщика.

Для выборки размером в 100-1000 элементов мы не замечали разницы в скорости обучения, поскольку из-за частой пересылки данных с/на GPU преимущество по скорости выполнения исчезало. Возможно, стоит организовать обучающий цикл так, чтобы данные оставались на GPU на все время обучения.