今回は、JuliaからGPUを利用してみたいと思います。
前回に引き続き、手元の環境を汚さないように、Dockerを利用して環境準備で行います。
前提
前提として、下記のソフトがインストールされていることを確認してください。
- 利用されているGPUに合わせたnvidia-driver
- Docker
Julia+CUDAの環境構築
以下の内容でDockerfile
を準備します。
FROM nvidia/cuda:11.4.2-cudnn8-runtime-ubuntu20.04
RUN apt-get update && apt-get upgrade -y
# julia
RUN apt-get install -y python3-pip
RUN pip3 install jill -U
RUN echo "Y" | jill install 1.6.2
RUN apt-get install -y curl wget
RUN julia -e 'using Pkg; Pkg.add(["Pluto"])'
RUN julia -e 'using Pkg; Pkg.add(["Flux", "CUDA"])'
私がM1 macを使っていることもあり、プラットフォームに依存させたくなかったので、JILL.pyを使ってJuliaのインストールを行いました。
Ubuntuを利用していれば、apt-get install -y julia
でもインストールは行うことができます。
次に下記のコマンドを利用して、作成したDockerfile
をビルドします。
ここで、julia_gpu
は作成するコンテナイメージの名前なので、自由に変更ができます。
$ docker build -t julia_gpu .
最後に下記コマンドでコンテナを立ち上げ、Plutoを起動します。
この際に--gpus
オプションで、DockerコンテナからGPUにアクセスできるようにします。
$ docker run -it --rm \
--gpus 1 \
-p 1234:1234 \
-v $(pwd):/workdir \
-w /workdir \
${IMAGE_NAME} \
/usr/local/bin/julia -e 'using Pluto; Pluto.run(host="0.0.0.0")'
以下のようなメッセージが表示されるので、http://0.0.0.0:1234/?secret=xxxxxx にブラウザからアクセスします。
Go to http://0.0.0.0:1234/?secret=xxxxxx in your browser to start writing ~ have fun!
Press Ctrl+C in this terminal to stop Pluto
CUDAのサンプルプログラムを実行
Pultoにアクセスしたら、「Create a new notebook」をクリックして、新しいノートブックを立ち上げます。
ノートブックのセルに以下のように入力します。
begin
using CUDA
W = cu(rand(2, 5)) # a 2×5 CuArray
b = cu(rand(2))
predict(x) = W*x .+ b
loss(x, y) = sum((predict(x) .- y).^2)
x, y = cu(rand(5)), cu(rand(2)) # Dummy data
loss(x, y)
end
入力したら、「Shift+Enter」でセルを実行します。
最初に実行する際は、わりと時間がかかると思います。
実行が完了すると、以下のようにセルの上側に実行結果(loss(x, y)
の計算結果)が表示されます。
Flux+CUDAでGPUを使ってディープラーニング
ついでに、Juliaで書かれたディープラーニングライブラリ「Flux.jl」でCUDAを利用してみましょう。
まずは、1つ目のセルでFlux
とCUDA
を読み込みます。
using Flux, CUDA
次のセルには、以下のように入力します。
begin
m = Dense(10,5) |> gpu
x = rand(10) |> gpu
m(x) |> cpu
end
Dense(10, 5)
は、入力10で出力5のFC層を示しています。
|> gpu
をつけることで、宣言した変数や関数をGPUに載せることができます。
結果を表示するには、CPUに戻す必要があるので|> cpu
を使います。
これを実行すると、以下のように5次元のベクトルが出力されます。
まとめ
いかがでしたでしょうか?
今回は、JuliaからGPUを利用したプログラミングを行ってみました。
本ブログでは、ローカル環境を守るために極力Dockerを利用していますが、ご自身の環境に直接インストールする際も同じ方法で行うことができます。
JuliaとGPUで高度な計算をガリガリ行いましょう!