logo

ジュリア・フラックスでGPUを使用する方法 📂機械学習

ジュリア・フラックスでGPUを使用する方法

概要

Juliaの機械学習ライブラリーであるFlux.jl1を使ってディープラーニングを実装する方法と、GPUを使って学習のパフォーマンスを加速する方法について紹介する。

GPUを使うためにはCUDA.jl2を使う必要があるし、前もってCUDAに関連する設定がされていなければならない。CUDAの設定自体はPythonでのそれと似ているので、次のポストを参考にしよう：

はっきりと言うが、GPUを使ったディープラーニングは基本でありながら、初心者にとって大きな壁でもある。その点で、Juliaの公式チュートリアルがあまり親切でない感じがするが、見方を変えれば、GPUの加速が非常に直感的で簡単に可能である。

コード

julia> using Flux, BenchmarkTools
       optimizer = Adam()
       
       f(v) = v[1].^2 + 2v[2] - v[3]*v[4]
       n = 10_000
       
       function ftry_FFNN()
           FFNN = Chain(
               Dense( 4, 100, relu),
               Dense(100, 100, relu),
               Dense(100, 100, relu),
               Dense(100, 1),
           )
           return FFNN
       end
ftry_FFNN (generic function with 1 method)

このポストでは、非線形関数を見つける問題として、順伝播ニューラルネットワークfeed-Forward Neural Network, FFNNのMLPを実装する例を取り上げる。(../3227) すぐに人工ニューラルネットワークのパフォーマンスが問題ではなく、GPUを使ってネットワークを学習する方法に焦点を当てよう。ftry_FFNN()という関数の名前は、FFNNを返す工場factoryという意味で決められたものだ。

julia> X = rand(Float32, 4, n)
4×10000 Matrix{Float32}:
 0.669836  0.260559  0.710337  0.5121    …  0.298401  0.763405  0.977941   0.89907
 0.135182  0.938298  0.110935  0.534417     0.804302  0.197353  0.784419   0.179469
 0.287801  0.839834  0.718759  0.15976      0.442106  0.696369  0.41352    0.617924
 0.213884  0.345568  0.426718  0.961369     0.542308  0.378965  0.0633196  0.580489

julia> Y = f.(eachcol(X))'
1×10000 adjoint(::Vector{Float32}) with eltype Float32:
 0.657488  1.65427  0.419741  1.17749  …  1.45789  0.713595  2.49902  0.808568

データはただランダムに生成する。

CPU

julia> data = Flux.DataLoader((X, Y), shuffle = true, batchsize = 1000);

julia> cFFNN = ftry_FFNN()
Chain(
  Dense(4 => 100, relu),                # 500 parameters
  Dense(100 => 100, relu),              # 10_100 parameters
  Dense(100 => 100, relu),              # 10_100 parameters
  Dense(100 => 1),                      # 101 parameters
)                   # Total: 8 arrays, 20_801 parameters, 81.754 KiB.

julia> Loss(x,y) = Flux.mse(cFFNN(x), y)
Loss (generic function with 1 method)

julia> ps = Flux.params(cFFNN);

julia> @btime Flux.train!(Loss, ps, data, optimizer)
  22.027 ms (3413 allocations: 60.07 MiB)

cpu.webp

エポックが繰り返されるごとに、CPUの占有率が上がるのが分かる。

GPU

julia> using CUDA

julia> CUDA.functional()
true

まずCUDA.jlをロードし、CUDA.functional()を通じてCUDAがうまく動作する環境か確認する作業が必要だ。trueが出力されれば、CUDAが動作しているという意味だ。次にすることは、たった２つだ。

一つ目、データをGPUに上げる：

julia> X = X |> gpu
4×10000 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 0.490089  0.305455  0.334825  0.0466784  …  0.243163  0.732825    0.401764  0.361136      
 0.102807  0.755049  0.971202  0.522919      0.766326  0.498049    0.669154  0.546359      
 0.440698  0.777547  0.263636  0.448606      0.854045  0.465477    0.314271  0.854085      
 0.288417  0.932674  0.808397  0.180423      0.63759   0.00347775  0.565556  0.872233      

julia> Y = Y |> gpu
1×10000 adjoint(::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}) with eltype Float32:
 0.318696  0.878203  1.84139  0.967079  0.162922  …  1.04725  1.53151  1.32198  0.478175   

julia> data = Flux.DataLoader((X, Y), shuffle = true, batchsize = 1000);

二つ目、関数をGPUに上げる：

julia> gFFNN = ftry_FFNN() |> gpu
Chain(
  Dense(4 => 100, relu),                # 500 parameters
  Dense(100 => 100, relu),              # 10_100 parameters
  Dense(100 => 100, relu),              # 10_100 parameters
  Dense(100 => 1),                      # 101 parameters
)                   # Total: 8 arrays, 20_801 parameters, 1.109 KiB.

そして、CPUでしていたことと同様に、普通に学習を進めればいい：

julia> Loss(x,y) = Flux.mse(gFFNN(x), y)
Loss (generic function with 1 method)

julia> ps = Flux.params(gFFNN);

julia> @btime Flux.train!(Loss, ps, data, optimizer)
  8.588 ms (24313 allocations: 1.90 MiB)

要するに、以下の３行だけでGPUでディープラーニングが可能になるということだ：

  • X |> gpu
  • Y |> gpu
  • FFNN |> gpu

gpu.webp

初心者がGPUを使うとき、最も重要なのは常にハードウェアのリソース消費量をモニタリングして、本当にデバイスを使っているか、CPUより遅くなっていないかを確認することだ。CUDAに関連する設定を間違えたり、GPUプログラミングに慣れていない場合、CPUだけで動いているのに「動いているから問題ないだろう」と安心しがちだ。

パフォーマンス比較

julia> @btime Flux.train!(Loss, ps, data, optimizer) # CPU
  23.120 ms (3413 allocations: 60.07 MiB)

julia> @btime Flux.train!(Loss, ps, data, optimizer) # GPU
  8.427 ms (24313 allocations: 1.90 MiB)

GPUの場合、CPUより約３倍速くなったのが見られるが、実際は「GPUを使って損をしていない」という程度が確認できるだけで、実環境ではもっと大きな差が出る。特に、データのサイズやアーキテクチャが大きくなるほど、そのようになる。筆者が研究で使っていたニューラルネットワークは、約50倍以上のパフォーマンス差を経験した。

一方で、パフォーマンスではなく「実装」自体に意味がある研究で、複雑な関数が使われている場合やインデクシングが必要な場合など、むしろGPUを使って損をする可能性がある。その場合は、CPUに戻すか、あるいはGPUで快適に動くようにネットワークを構成するために色々なトリックを使うこともできるが、GPUで快適に動くようにレイヤーを構成することは、個々の能力に大きく依存する。

全コード

using Flux

optimizer = Adam()

f(v) = v[1].^2 + 2v[2] - v[3]*v[4]
n = 10_000

function ftry_FFNN()
    FFNN = Chain(
        Dense( 4, 100, relu),
        Dense(100, 100, relu),
        Dense(100, 100, relu),
        Dense(100, 1),
    )
    return FFNN
end

X = rand(Float32, 4, n)
Y = f.(eachcol(X))'

data = Flux.DataLoader((X, Y), shuffle = true, batchsize = 1000);

cFFNN = ftry_FFNN()

Loss(x,y) = Flux.mse(cFFNN(x), y)
ps = Flux.params(cFFNN);
@btime Flux.train!(Loss, ps, data, optimizer)

# --------------------

using CUDA

CUDA.functional()
X = X |> gpu
Y = Y |> gpu
data = Flux.DataLoader((X, Y), shuffle = true, batchsize = 1000);
gFFNN = ftry_FFNN() |> gpu

Loss(x,y) = Flux.mse(gFFNN(x), y)
ps = Flux.params(gFFNN);
@btime Flux.train!(Loss, ps, data, optimizer)

using BenchmarkTools

@btime Flux.train!(Loss, ps, data, optimizer) # CPU
@btime Flux.train!(Loss, ps, data, optimizer) # GPU

環境

  • OS: Windows
  • julia: v1.9.0
  • Flux v0.13.16

  1. https://github.com/FluxML/Flux.jl ↩︎

  2. https://github.com/JuliaGPU/CUDA.jl ↩︎