すでに以下のようなわかりやすい記事がありますが、流石に月日が流れ、PyCall.jlやPyTorchの方にも仕様変更があったようなので、少し自分でも見てみようと思いました。
(個人的にFlux.jl使ってたんですが、ちょっとZygote.jlではどうにもならない自動微分の案件がありまして......)
↑こういう経緯なので、基本自分用のメモみたいなもので、個人的に必要ないと感じたことは飛ばしますので、ご容赦を。
環境構築は本家https://github.com/JuliaPy/PyCall.jl を見たほうが良いです。僕はすでに入れてあったAnaconda3のPyTorchが入った仮想環境を使っています。
PyCall.jlのキホン
pythonのライブラリはpyimport()で出来ます。
using PyCall
torch = pyimport("torch")
ここで定義した変数torchはjuliaではPyObjectという構造体になっています。
pythonライブラリ内で定義されたクラスも
nn = torch.nn
F = torch.nn.functional
のように使う事ができます1。これらもPyObjectです。
PyCall.jlでは、一部の代表的な型、例えばjuliaのArrayとpythonのnumpy.ndarrayを自動的に変換してくれたりするのですが、それ以外は一括でPyObjectで扱っています。
例えば
A = torch.zeros((2, 3))
とすれば
PyObject tensor([[0., 0., 0.],
[0., 0., 0.]])
となります。
A.cuda()
とすればもちろんGPUにも転送できます。
PyObject tensor([[0., 0., 0.],
[0., 0., 0.]], device='cuda:0')
torch.tensorをjuliaで扱えるようにするにはnumpy.ndarrayに変換すれば、先程書いたようにjulia上では自動的にArrayとして変換されます。
A.numpy()
2×3 Matrix{Float32}:
0.0 0.0 0.0
0.0 0.0 0.0
ここまではかなり直感的だと思います。
pythonはオブジェクト指向なのでクラスも扱える必要があります。ここは少しだけ違っていて、
@pydef mutable struct Test
function __init__(self)
self.x = 0
end
function add1(self)
self.x += 1
end
end
のようにmutable structに@pydefマクロをつける2ことでpythonのクラスに対応するものを定義できます。メソッドは関数として定義すればOKです。一応対応するpythonコードも示しておきます。
class Test:
def __init__(self):
self.x = 0
def add1(self):
self.x += 1
実際に確かめると、
test = Test()
println(test.x)
test.add1()
println(test.x)
0
1
きちんと動いています。
ちなみに継承は<:を使います。
@pydef mutable struct Test2 <: Test
function __init__(self)
pybuiltin("super")(Test2, self).__init__()
# pybuiltin(:super)(Test2, self).__init__() でもOK
end
function mult3(self)
self.x *= 3
end
end
class Test2(Test):
def __init__(self):
super().__init__()
# same as super(Test2, self).__init__()
def mult3(self):
self.x *= 3
ここでpythonのsuper().__init__() がpybuiltin("super")(Test2, self).__init__()に変わっています。
pybuiltin()は引数の文字列(もしくはシンボル)に対応するpythonの組み込み関数を呼び出す関数です。
pythonのsuper()は引数を省略するとよしなにやってくれるのですが、pybuiltin()は引数を明示しないといけないようです。
test2 = Test2()
println(test2.x)
test2.add1()
println(test2.x)
test2.mult3()
println(test2.x)
0
1
3
ニューラルネットワーク
本題に入ることにします。難しい例をやっても仕方ないので、ここではレナードジョーンズ型のポテンシャル関数 $U(r) = \epsilon[(\sigma/r)^{12}-(\sigma/r)^6]$をニューラルネットワークで回帰することにします。
上記の内容とPyTorchの扱いがわかっていればほとんど直感的にできると思いますので、説明はあっさり風味にします。
using PyCall
using Plots
using Zygote
torch = pyimport("torch")
nn = torch.nn
F = torch.nn.functional
device = torch.device(ifelse(torch.cuda.is_available(), "cuda:0", "cpu"))
PyObject device(type='cuda', index=0)
続いてレナードジョーンズポテンシャルを定義します。
データはポテンシャル$U(r)$だけでなくその勾配の符号反転である力$F(r) = -\frac{d}{dr}U(r)$も含めることにします。
力はZygote.jlの微分で出しています。
mutable struct LJ
σ
ϵ
end
(m::LJ)(r) = m.σ[1] .* (m.ϵ[1] ./ r).^12 - m.σ[1] .* (m.ϵ[1] ./ r).^6
σ_true = [1.0]
ϵ_true = [2.0]
lj = LJ(σ_true,ϵ_true)
lj_force(r) = gradient(r -> -lj(r)[1], r)[1]
(パラメータを配列にしたり、call overloadingしているのはFlux.jl使ってたときの名残なので気にしない)
データ作ってtorch.tensorにしてgpuに流します。
r_train = reshape(range(2.0,5.0,31),:,1)
U_train = lj(r_train)
F_train = lj_force.(r_train)
r_train = torch.tensor(r_train,requires_grad=true,dtype=torch.float32).to(device)
F_train = torch.tensor(F_train,dtype=torch.float32).to(device)
U_train = torch.tensor(U_train,dtype=torch.float32).to(device)
ニューラルネットワークの定義をします。
@pydef mutable struct Net <: torch.nn.Module
function __init__(self)
pybuiltin("super")(Net, self).__init__()
self.net1 = torch.nn.Linear(1, 5)
self.net2 = torch.nn.Linear(5, 1)
end
function forward(self, x)
x = self.net1(x)
x = torch.sigmoid(x)
x = self.net2(x)
return x
end
end
学習します。損失関数として、ポテンシャルと力の二乗誤差の混合を使います。
net = Net()
net.to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-2, betas=(0.9, 0.99), eps=1e-07)
loss_list = []
epoch = 2000
for i = 1:epoch
U_pred = net(r_train)
F_pred = []
for i = 1:U_pred.shape[1]
push!(F_pred, - torch.autograd.grad(U_pred[i], r_train, create_graph=true)[1][i])
end
F_pred = torch.cat(F_pred).reshape(-1, 1)
loss = criterion(F_pred, F_train) + 10 * criterion(U_pred, U_train)
loss_item = loss.item()
println("Epoch: $i, Loss: $loss_item")
push!(loss_list,loss_item)
optimizer.zero_grad()
loss.backward()
optimizer.step()
end
学習したモデルを示します。
r_pred = reshape(range(2.0,5.0,301),:,1)
r_pred = torch.tensor(r_pred,requires_grad=true,dtype=torch.float32).to(device)
U_pred = net(r_pred)
F_pred = zeros(U_pred.shape[1])
for i = 1:U_pred.shape[1]
F_pred[i,1] = - torch.autograd.grad(U_pred[i], r_pred, create_graph=true)[1][i]
end
U_pred = U_pred.cpu().detach().numpy()
r_pred = r_pred.cpu().detach().numpy()
U_train = U_train.cpu().detach().numpy()
r_train = r_train.cpu().detach().numpy()
F_train = F_train.cpu().detach().numpy()
表示してみます。
scatter(r_train[:], U_train[:],legend=:none)
plot!(r_pred[:], U_pred[:],legend=:none,title="potential")
scatter(r_train[:], F_train[:],legend=:none)
plot!(r_pred[:], F_pred[:],legend=:none,title="force")
よく一致していますね。汎化性能については検証していないですが、学習は出来ていそうです。
追記：DataLoaderの使い方
pythonで扱う場合はDataLoaderは楽に使えるのですが、PyCall.jlを使って呼び出す場合は少し癖があります。
python形式のDataLoaderはjuliaから見ると謎のPyObjectでしかないので、for文で回せないからです。
なので原始的にDataLoaderをiter()でiteratorにしてからnext()で進める、という操作をループさせます3。最後まで行った状態でさらにnext()するとエラーとなるので、try~~catch構文でループを止めます。
torch_data_set = data.TensorDataset(r_train, U_train, F_train)
loader = data.DataLoader(dataset=torch_data_set, batch_size=10,shuffle=true,drop_last=true)
loader_iter = pybuiltin("iter")(loader)
while true
try
r_batch, U_batch, F_batch = pybuiltin("next")(loader_iter)
### ミニバッチへの処理 ###
catch
break
end
end
という感じです。
わかってしまえば難しくはないですね。
一応DataLoader使ったバージョンも記載しておきます(このレベルのデータ数、次元、関数の複雑さでは、ミニバッチSGDは最適化の単なるノイズでしかないので性能は落ちます)。
# dataloader
r_train = reshape(range(2.0,5.0,151),:,1)
U_train = lj(r_train)
F_train = lj_force.(r_train)
r_train = torch.tensor(r_train,requires_grad=true,dtype=torch.float32)
F_train = torch.tensor(F_train,dtype=torch.float32)
U_train = torch.tensor(U_train,dtype=torch.float32)
# r_train = torch.tensor(r_train,requires_grad=true,dtype=torch.float32).to(device)
# F_train = torch.tensor(F_train,dtype=torch.float32).to(device)
# U_train = torch.tensor(U_train,dtype=torch.float32).to(device)
torch_data_set = data.TensorDataset(r_train, U_train, F_train)
loader = data.DataLoader(dataset=torch_data_set, batch_size=30, shuffle=true,drop_last=true)
@pydef mutable struct Net <: torch.nn.Module
function __init__(self)
pybuiltin("super")(Net, self).__init__()
self.net1 = torch.nn.Linear(1, 5)
self.net2 = torch.nn.Linear(5, 1)
end
function forward(self, x)
x = self.net1(x)
x = torch.sigmoid(x)
x = self.net2(x)
return x
end
end
net = Net()
net.to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-2, betas=(0.9, 0.99), eps=1e-07)
loss_list = []
epoch = 10000
for i = 1:epoch
loader_iter = pybuiltin("iter")(loader)
batch_i = 0
loss_item = 0.0
while true
try
r_batch, U_batch, F_batch = pybuiltin("next")(loader_iter)
r_batch = r_batch.requires_grad_(true).to(device)
U_batch = U_batch.to(device)
F_batch = F_batch.to(device)
U_pred = net(r_batch)
F_pred = []
for i = 1:U_pred.shape[1]
push!(F_pred, - torch.autograd.grad(U_pred[i], r_batch, create_graph=true)[1][i])
end
F_pred = torch.cat(F_pred).reshape(-1, 1)
loss = criterion(F_pred, F_batch) + 10 * criterion(U_pred, U_batch)
loss_item += loss.item()
batch_i += 1
optimizer.zero_grad()
loss.backward()
optimizer.step()
catch
break
end
end
loss_item /= batch_i
println("Epoch: $i, Loss: $loss_item")
push!(loss_list,loss_item)
end