Engee 文档
Notebook

CPU 和 GPU 神经网络训练的比较

让我们回到为回归任务训练多层神经网络的例子,比较一下在切换到 GPU 时将如何加速这一过程。

安装程序库

如果尚未安装指定的程序库,执行以下命令可能需要几分钟时间。

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

准备工作

让我们创建相当大的训练数据向量,如Float32

In [ ]:
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);

并设置足够长的训练时间:

In [ ]:
epochs = 2000;
learning_rate = 0.08;

提升 CPU 性能时,最好从epochs=100 开始,然后慢慢增加。即使我们离 "工业 "值还很远,但在次优设置下训练神经网络时,我们仍有可能损失大量 CPU 时间。

CPU 上的数据加载和训练

让我们创建一个标准训练程序:

In [ ]:
model_cpu = Chain( Dense( 2 => 20, relu ), Dense( 20 => 5, relu ), Dense( 5 => 1 ) )
model_copy = deepcopy( model_cpu ) # На будущее нам будет нужна точно такая же модель
data = [ (Xs', Ys') ]
opt_state = Flux.setup( Adam( learning_rate ), model_cpu );
In [ ]:
using Random
Random.seed!( 1 ); # Установим генератор случайных чисел в воспроизводимое состояние
In [ ]:
@time (for i in 1:epochs
    Flux.train!(model_cpu, data, opt_state) do m, x, y
        Flux.mse( m(x), y ) # Лосс функция
    end
end)
467.851443 seconds (20.01 M allocations: 27.787 GiB, 1.23% gc time, 3.11% compilation time)

训练需要很长时间。如果我们的工作受到 CPU 的限制,我们就必须减少训练样本中的点数,至少在训练的第一阶段,我们要选择神经网络的超参数。

其中 1-2% 的时间用于清理内存(gc_time = garbage collection )和编译代码(compile time )。

将进程传输到 GPU

通过连接Flux 库,我们可以使用|> gpu 结构(将左侧表达式传递给函数 gpu())。它允许我们将矩阵或结构发送至 GPU 并在其上执行代码,而无需创建额外的嵌套数据处理层。

In [ ]:
if CUDA.functional()
    model_gpu = model_copy |> gpu;
    
    Xg = Xs' |> gpu;
    Yg = Ys' |> gpu;
    data = [ (Xg, Yg) ]
    
    opt_state = Flux.setup( Adam( learning_rate ), model_gpu );
end;

请注意,我们曾尝试使用DataLoader 。它的使用减慢了在 GPU 上的执行速度,由此我们得出结论,我们应该相信 Julia GPU 编译器实现的并行化程序。

In [ ]:
Random.seed!( 1 ); # Установим генератор случайных чисел в воспроизводимое состояние
In [ ]:
if CUDA.functional()
    @time CUDA.@sync (
        for i in 1:epochs
            Flux.train!(model_gpu, data, opt_state) do m, x, y
                Flux.mse( m(x), y ) # Лосс функция
            end
        end
    )
end
 26.461145 seconds (33.51 M allocations: 1.926 GiB, 2.39% gc time, 80.03% compilation time: 2% of which was recompilation)

GPU 移植总结

一般来说,将神经网络移植到 GPU 需要三个额外步骤:

  • 传输特征矩阵Xs |> gpu
  • 传输响应矩阵Ys |> gpu
  • 传输神经网络结构model |> gpu

对于少量训练数据(100-200 个示例),CPU 的神经网络学习速度更快。但在 40000 个示例时,GPU 的优势就显现出来了。

如果我们减去编译时间(compile time )和执行后的内存清理时间(gc time = * 垃圾收集时间*),我们可以说

**显卡比 CPU 快约 200 倍。

检查结果

我们使用 GPU 训练的神经网络对原始数据集进行插值:

In [ ]:
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)
)
Out[0]:

有趣的是,只需设置Random.seed 就能获得两个完全相同(看似相同)的结果--一个在 CPU 上,一个在 GPU 上。

结论

我们在 CPU 和 GPU 上训练了一个神经网络来插值一个包含 40,000 个元素的样本。

在 GPU 上进行 2,000 次历时训练的速度比 CPU 快近 200 倍,因此可以进行更多的超参数迭代和训练更大的网络,从而节省了设计者的时间。

对于 100-1000 个元素的样本量,我们没有注意到训练速度上的任何差异,因为从 GPU 到 GPU 的频繁数据传输导致速度优势消失。也许我们应该以这样一种方式来安排训练周期,即数据在整个训练过程中都保留在 GPU 上。