Pythonいぬ

pythonを使った画像処理に関する記事を書いていきます

Pytorch で線形回帰 (Linear Regression) を実装してみる

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()

結果の図。青がデータで赤が最適化で得られた直線。

f:id:tzmi:20200115231250p:plain