昨年のアドベントカレンダーで、PyCall.jlでPyTorch使ってDeep Learningするという投稿をしましたが、その続きです。（なんか、１年前と同じことやってますね。。）
昨年は、PyCall.jlを使ってJuliaからPyTorchを呼ぶという話でしたが、今年は、逆に、PyTorch（Python）からpyjuliaを使ってJuliaを呼ぶ、という話です。
愚痴
自分１人だけなのでであれば昨年書いたように、JuliaからPyTorchを呼ぶでいいんですが、いかんせん複数人でいろんなモデルを作って試すとなると、みんなPython使ってるわけで、やっぱり、現状ではPythonをメインに使わざるをえない部分があります。
実際、PyTorch便利ですよね。なんだかんだいって、Flux.jlやKnet.jlなどとは、フレームワークとしての完成度が違います。
なんですけど、PyTorch（というかPython）なんかで、深層学習なり数値計算するとなると、とにかくベクトル演算で一度に処理するように書かないといけない（For文は使ってはいけない）という強い縛りがあるわけです。なにしろ、PythonのFor文は死ぬほど遅いですからね。
もちろん、普通の、CNNとかLSTMとかなら、PyTorchが用意しているLayerを呼べばいいだけなんですが、いかんせん、ちょっと複雑なことをする独自Layerを作りたいとかなると、（PyTorchに用意されている関数では）どうやってもベクトル演算として書けないなんて処理もでてきます。
そういう場合、普通は、C++あるいはCUDAでextentionを作るということになるわけですが、C++ではなくてJuliaで書きたい。なにしろ、Juliaなら（GPUArrays.jl使えば）、CPUだけではなくて、GPUも簡単に使えるんですよ。
だいたい、そうでなくても、いったん Juliaに一度なれてしまうと、とにかく、Python（numpy、PyTorch）の全てをベクトル演算で書け（For文を使ったら負け）という縛りは、ほんとにキツイです。やってられないです。もしかして、私だけ？そんなことないですよね、みんなそう思いますよね。
人間の自然な思考はFor文であって、ベクトル演算ではないですよね。
愚痴がすぎました。
pyjulia使って、PyTorchのTensorをjuliaでいじる
というわけで本題です。
pyjuliaを使うと、PyTorchのTensorをJuliaでいじれます。C/C++などでグルーコードを書く必要もありません。pythonのスクリプトの中にJuliaのスクリプトを文字列として埋め込むだけです。（もちろん、PythonとJuliaを別ファイルに分けることもできます）
PyTorchのカスタムレイヤーをJuliaで書く、それも、CPUにもGPUも同一のコードで対応できます。device agnosticというやつですね。
あるいは、カスタムレイヤーだけではなくて、結果の可視化なんかでも、ちょっと独自の可視化をやりたいとかなると、Pythonだとかなりつらかったりするんですが、Juliaの描画ライブラリ業界における期待の新星Makie.jlなんかを使えば、大量のニューロンの状態を独自の方法で描画したい、とかも高速にできます。
pyjuliaのインストール
pyjuliaの公式ページにしたがって、pyjuliaをインストールします。
注意点としては、公式ページにも書いてありますが、anaconda版のPythonではそのままでは動きません。
python-jlを使うか、あるいは、pyenvで--enable-sharedオプションを使ってPythonをインストールする必要があります。
私は、anaconda版のPythonをインストールしたあと、pyenvに含まれるpython-buildを使って、--enable-sharedオプション版のPythonをanaconda版のpythonの上に上書きしてしまう、という強引な方法を使っています。
anaconda版のPythonとpython-buildするPythonのバージョンを合わせれば、とりあえずは、問題なく動きます。
PyTorchからJuliaを呼ぶ
まずは、以下の内容のJuliaスクリプトを、カレントディレクトリにPyTorchJulia.jlという名前でつくります。
（本当は、真面目にPackageとして作るのがいいんでしょうけど、面倒なので）
これが、PyTorchのTensorをJuliaのArrayやCuArrayに変換するインタフェースのコードです。
module PyTorchJulia
export jl_array, jl_type
using GPUArrays
using CuArrays
using PyCall
const torch = pyimport("torch")
const is_cuarray_unsafe_wrap = :own in fieldnames(CuArray)
const py_str = pybuiltin("str")
const torch_types = Dict(
"torch.float32" => Float32,
"torch.float64" => Float64,
"torch.float16" => Float16,
"torch.uint8" => UInt8,
"torch.int8" => Int8,
"torch.int16" => Int16,
"torch.int32" => Int32,
"torch.int64" => Int64)
jl_type(dtype) = torch_types[pycall(py_str, String, dtype)]
function jl_array(x::PyObject)
@assert x[:is_contiguous]()
T = jl_type(x[:dtype])
dims = reverse(x[:shape])
device = x[:device]
device_type = device[:type]
if device_type == "cpu"
ptr = Ptr{T}(x["data_ptr"]())
return unsafe_wrap(Array, ptr, dims)
elseif device_type == "cuda"
if is_cuarray_unsafe_wrap
ptr = Ptr{T}(x["data_ptr"]())
return unsafe_wrap(CuArray, ptr, dims)
else
ctx = CuArrays.CuCurrentContext()
@assert CuArrays.CUDAdrv.device(ctx) == CuArrays.CuDevice(device[:index])
buf = CuArrays.Mem.Buffer(
Ptr{Cvoid}(x["data_ptr"]()),
x[:numel]() * sizeof(T),
ctx)
CuArrays.Mem.retain(buf)
return CuArray{T, length(dims)}(buf, dims)
end
end
end
GPUArrays.backend(::Type{<:Array}) = GPUArrays.JLBackend()
end
で、これを使う、Python（PyTorch）のコードを書きます。
import torch
from julia import Main as julia
jl_script = '''
include("PyTorchJulia.jl")
using .PyTorchJulia
using GPUArrays
using PyCall
const torch = pyimport("torch")
Base.getproperty(o::PyObject, s::Symbol) = isdefined(o, s) ? getfield(o, s) : getindex(o, s)
f_impl!(x) = (x .*= 2)
f!(x) = f_impl!(jl_array(x))
g_impl!(out, x) = (@. out = x * 2)
function g(x)
out = torch.empty_like(x)
g_impl!(jl_array(out), jl_array(x))
return out
end
function h_impl!(out, x, y)
gpu_call(out, (out, x, convert(eltype(x), y))) do state, out, x, y
idx = @linearidx(out)
out[idx] = x[idx] + y
return
end
end
function h(x, y)
out = torch.empty_like(x)
h_impl!(jl_array(out), jl_array(x), y)
return out
end
'''
julia.eval(jl_script)
x_org = torch.empty((2,3,4), dtype=torch.float32).uniform_()
x_cpu = x_org.clone()
julia.f_b(x_cpu)
assert torch.allclose(x_cpu, 2 * x_org)
x2_cpu = julia.g(x_cpu)
assert torch.allclose(x2_cpu, 2 * x_cpu)
x3_cpu = julia.h(x_cpu, 3.0)
assert torch.allclose(x3_cpu, x_cpu + 3)
if torch.cuda.is_available():
x_gpu = x_org.cuda()
julia.f_b(x_gpu)
assert torch.allclose(x_gpu.cpu(), x_cpu)
x2_gpu = julia.g(x_gpu)
assert torch.allclose(x2_gpu.cpu(), x2_cpu)
x3_gpu = julia.h(x_gpu, 3.0)
assert torch.allclose(x3_gpu.cpu(), x3_cpu)
Pythonのスクリプト上で文字列として定義したJuliaの関数（f!(x), g(x), h(x, y)）を、
PyTorchのTensorであるx_cpuを引数として
それぞれ、julia.f_b(x_cpu)，x2_cpu = julia.g(x_cpu)，x3_cpu = julia.h(x_cpu, 3.0) と呼び出しています。
あるいは、PyTorchのCUDA Tensorである、x_gpuを引数として、
それぞれ、julia.f_b(x_gpu)，x2_gpu = julia.g(x_gpu)，x3_gpu = julia.h(x_gpu, 3.0) と呼び出しています。
簡単な解説
-
Pythonではビックリマーク
!を識別名として使えないため、juliaのf!という名前の関数は、pyjuliaの仕様として、ビックリマーク!が_bに置換されて、Pythonからはf_bとなります。 -
PyTorchJulia.jlで定義しているjl_array()という関数で、PyTorchのTensorをJuliaのArrayあるいはCuArrayに変換しています。
PyTorchのTensorのメモリをそのままJuliaのArrayやCuArrayとしてラップしているだけで、コピーすることなく高速に変換できます。 -
ただし、PyTorchのTensorのメモリ配置がRow-mejorであるのに対して、JuliaのArrray（やCuArray）はColumn-majorです。jl_arrayではメモリのコピーをしないかわりに、PyTorchのTensorをJuliaのArray（やCuArray）に変換すると、次元の順番が逆になります。例えば、PyTorchの次元(2,3,4)のTensorを、jl_arrayでJuliaのArrayに変換すると、次元は(4,3,2)になります。
（Arrayを使う代わりに、例えば、Strided.jlを使えば、次元をひっくり返さないですみそうです。もし、Package化するとなれば、対応したいです） -
PyTorchのTensorがCPU上にある場合にはJuliaのArrayに、GPU上にある場合にはCuArrayに変換されます。上のコードでは、
f_impl!(x),g_impl(x, y),h_impl!(out, x, y)はArrayかCuArrayかに関わらず同一コードで動作します。（device agnostic)
とくに、h_impl!(out, x, y)では、GPUArrays.jlのgpu_callを使って、CUDAのカーネルをPure Juliaで記述して呼んでいます。 -
juliaのコードで、
f!(x)とf_impl!(x)を別の関数に分けている（一つの関数にまとめていない）のは、function barrierを使いたいためです。
jl_array()の戻り値は、PyTorchのTensorがCPUにあるかGPUにあるかによって、Array型かCuArray型になります。関数を分けることで、f_impl!(x)は、xの型に応じて特殊化されたコードがコンパイルされるので、型安定なコードになって、高速に計算できます。 -
Juliaコード内で、新たなTensorを確保してPyTorch側に返したい場合には、
g(x)の中で行っているようにout = torch.empty_like(x)という形で、PyTorchのメソッドを使ってPyTorch側でTensorを確保する必要があります。 -
本当は、昨年の記事で詳しく説明したように、juliaからPython（PyTorch）の関数を呼ぶ場合には、PyCall.jlの仕様により、
torch[:empty_like](x)と書かなければなりません。これは、とても面倒なので、v1.0で導入されたBase.getpropertyをオーバーロードすることで、out = torch.empty_like(x)と呼べるようにしています。
なお、PyCall.jl公式でも、この記法を可能にするPRが提案されています。 -
PyTorchJulia.jlの中で、is_cuarray_unsafe_wrapの値に応じて、unsafe_wrap(CuArray, ptr, dims)、あるいは、CuArray{T, length(dims)}(buf, dims)を呼んでいます。実は、この記事を書いている途中で、CuArrays.jlに、unsafe_wrapを作ってほしいと提案したところ、公式に対応していただきました。すでに、masterにマージされていますが、まだ、Releaseにはなっていないので、旧版にも対応できるようにif文でわけています。
JuliaからPyTorchを使う場合
Juliaをメインに使いたい場合、つまり、昨年度の記事のように、JuliaからPyCall.jl経由でPyTorchを使う場合は、以下のJuliaスクリプトのようになります。
Base.getpropertyをオーバーロードしたことで、PythonとJuliaのコードはほぼコピペで動くようになりました。
include("PyTorchJulia.jl")
using .PyTorchJulia
using GPUArrays
using PyCall
const torch = pyimport("torch")
Base.getproperty(o::PyObject, s::Symbol) = isdefined(o, s) ? getfield(o, s) : getindex(o, s)
f_impl!(x) = (x .*= 2)
f!(x) = f_impl!(jl_array(x))
g_impl!(out, x) = (@. out = x * 2)
function g(x)
out = torch.empty_like(x)
g_impl!(jl_array(out), jl_array(x))
return out
end
function h_impl!(out, x, y)
gpu_call(out, (out, x, convert(eltype(x), y))) do state, out, x, y
idx = @linearidx(out)
out[idx] = x[idx] + y
return
end
end
function h(x, y)
out = torch.empty_like(x)
h_impl!(jl_array(out), jl_array(x), y)
return out
end
x_org = torch.empty((2,3,4), dtype=torch.float32).uniform_()
x_cpu = x_org.clone()
f!(x_cpu)
@assert torch.allclose(x_cpu, 2 * x_org)
x2_cpu = g(x_cpu)
@assert torch.allclose(x2_cpu, 2 * x_cpu)
x3_cpu = h(x_cpu, 3.0)
@assert torch.allclose(x3_cpu, x_cpu + 3)
if torch.cuda.is_available()
x_gpu = x_org.cuda()
f!(x_gpu)
@assert torch.allclose(x_gpu.cpu(), x_cpu)
x2_gpu = g(x_gpu)
@assert torch.allclose(x2_gpu.cpu(), x2_cpu)
x3_gpu = h(x_gpu, 3.0)
@assert torch.allclose(x3_gpu.cpu(), x3_cpu)
end
今後の展望（野望）
JuliaからPyTorchを使う場合、PyCall.jlを使ってPython経由で使うということになるわけですが、せっかくJulia使っているんで、できればPythonのお世話になりたくないと思ったりします。
実は、PyTorchは、もうすぐReleaseされる予定のv1.0から、C++ (C++11以降)でフロントエンドを記述できるようになります。公式ページを見ていただくとわかりますが、C++のコードはPythonのコードとほぼ一対一に対応すると言ってよいです。
となると、このC++ Frontend APIを、最近Julia v1.0に向けて作業が進んでいるCxx.jlを使ってラップすれば、Juliaから直接PyTorchを使うことができることになりそうです。できたら、やってみたいところです。