JuliaでCUDAを利用する【Docker】

今回は、JuliaからGPUを利用してみたいと思います。
前回に引き続き、手元の環境を汚さないように、Dockerを利用して環境準備で行います。

前提
Julia+CUDAの環境構築
CUDAのサンプルプログラムを実行
Flux+CUDAでGPUを使ってディープラーニング
まとめ

前提

前提として、下記のソフトがインストールされていることを確認してください。

利用されている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で高度な計算をガリガリ行いましょう！