Pytorchを使って1次元の線形回帰 (Linear Regression) の学習を実装してみる。目的はPytorchのモデルの書き方や学習コードの書き方の再確認。
線形回帰モデル
まずはPytorchで線形回帰モデルを書いてみる。パラメータは傾き(weight)と切片(bias)の2つ
from torch import nn class LinearRegression(nn.Module): def __init__(self): super().__init__() # python3では引数省略可能 self.layer = nn.Linear(1, 1, bias=True) def forward(self, x): y = self.layer(x) return y
pytorchのモデルはクラスで定義すると簡単で、レイヤを定義するコンストラクタ init(self) と forward 処理 foward(self, x) で構成される。モデルが大きくなってもだいたいこの形で定義できる。コンストラクタにはnn.Moduleの継承部分 super().init() が含まれる。
なお、パラメータへのアクセスは以下のようにする。
a_ = self.layer.weight b_ = self.layer.bias
モデルの外からのアクセスは以下
a_ = model.layer.weight b_ = model.layer.bias
numpy arrayへ変換する場合
a_ = model.layer.weight.detach().to('cpu').numpy() b_ = model.layer.bias.detach().to('cpu').numpy()
学習部分
上記で定義したモデルを使って学習部分を書いてみる。線形回帰だけではなく、もっと複雑なニューラルネットモデルを学習させる場合もこの書き方でOKかと思う。
import numpy as np import torch from torch import optim if __name__ == '__main__': # GPUかCPUかを自動設定 device = 'cuda' if torch.cuda.is_available() else 'cpu' # modelとoptimizerの定義 model = LinearRegression().to(device) opt = optim.SGD(model.parameters(), lr=0.01) x, y = getdata_function() # この行はダミー n = 1000 # number of data bs = 10 # batch_size niter = 1000 # number of iteration for iiter in range(niter): # batch dataの取得 r = np.random.choice(n, bs, replace=False) bx = x[r].reshape(-1,1) by = y[r].reshape(-1,1) # forwardとloss計算 y_ = model.forward(bx) loss = torch.mean((y_ - by)**2) # 最適化 opt.zero_grad() # 勾配初期化 loss.backward() # 勾配計算(backward) opt.step() # パラメータ更新
model定義の部分では最低 model と optimizer だけ定義すればよい。明示的にforループでイテレーションを作って最適化をまわす。lossはnn.MSELossなどを使ってもよいが、普通のtorchの計算でも大丈夫。最適化部分はzero_grad(), backward(), step()の順に並べておけばよい。
サンプルコード
以下にサンプルコード全体を貼り付ける。
import numpy as np from matplotlib import pyplot as plt import torch from torch import nn from torch import optim # 線形回帰モデルの定義 class LinearRegression(nn.Module): def __init__(self): super().__init__() self.layer = nn.Linear(1, 1, bias=True) def forward(self, x): y = self.layer(x) return y # ここからmain if __name__ == '__main__': # GPUかCPUかを自動設定 device = 'cuda' if torch.cuda.is_available() else 'cpu' # modelとoptimizerの定義 model = LinearRegression().to(device) opt = optim.SGD(model.parameters(), lr=0.01) # data生成 n = 1000 x = torch.rand(n)*2-1 a, b = 2.0, -10.0 # weight & bias y = a*x+b # dataにノイズ追加 x = x + torch.randn(n)*0.02 y = y + a*torch.randn(n)*0.02 # to GPU x = x.to(device) y = y.to(device) bs = 10 # batch_size niter = 1000 losses = [] for iiter in range(niter): # batch dataの取得 r = np.random.choice(n, bs, replace=False) bx = x[r].reshape(-1,1) by = y[r].reshape(-1,1) # forwardとloss計算 y_ = model.forward(bx) loss = torch.mean((y_ - by)**2) # 最適化 opt.zero_grad() # 勾配初期化 loss.backward() # 勾配計算(backward) opt.step() # パラメータ更新 print('%05d/%05d loss=%.5f' % (iiter, niter, loss.item())) losses.append(loss.item()) # 重みの取り出し a_ = model.layer.weight.detach().to('cpu').numpy() b_ = model.layer.bias.detach().to('cpu').numpy() print('a=%.3f b=%.3f' % (a_[0] ,b_[0])) # データと最適化した関数のplot xnp = x.detach().to('cpu').numpy() ynp = y.detach().to('cpu').numpy() plt.scatter(xnp, ynp) x = np.linspace(-1,1,100) y = a_[0]*x + b_[0] plt.plot(x, y, c='r') plt.savefig('output.png') plt.tight_layout();plt.show()
結果の図。青がデータで赤が最適化で得られた直線。