ローリングコンバットピッチなう!

AIとか仮想化とかペーパークラフトとか

【PyTorch始めました】Chainerで書いたディープラーニングのコードをPyTorchに移植する(とりあえずMNIST)

[technology][PyTorch]対訳MNIST: Chainer vs. PyTorch


無事にPyTorchを使える環境が構築できたので、Chainer用に組んだスクリプトをPyTorchに移行させてみます。
題材はお馴染みのMNISTです。
Chainerのコードは以下の通り。

# -*- Coding: utf-8 -*-

# Numpy
import numpy as np
# Chainer
import chainer
import chainer.links as L
import chainer.functions as F
from chainer import Chain,optimizers,Variable

train, test = chainer.datasets.get_mnist()

x_train = []
x_test  = []
t_train = []
t_test  = []

for line in train:
    x_train.append(line[0])
    t_train.append(int(line[-1]))

for line in test:
    x_test.append(line[0])
    t_test.append(int(line[-1]))

x_train = np.array(x_train,dtype=np.float32)
x_test = np.array(x_test,dtype=np.float32)
t_train = np.array(t_train,dtype=np.int32)
t_test = np.array(t_test,dtype=np.int32)

# DNN class
class DNN(Chain):
    def __init__(self):
        super(DNN, self).__init__(
            l1 = L.Linear(None,1000),
            l2 = L.Linear(None,1000),
            l3 = L.Linear(None,10)
        )
    def forward(self,x):
        h1 = F.relu(self.l1(x))
        h2 = F.relu(self.l2(h1))
        h3 = self.l3(h2)
        return h3

# Create DNN class instance
model = DNN()

# Set optimizer
optimizer = optimizers.Adam()
optimizer.setup(model)

# Number of epochs
n_epoch = 10

# batch size
batch_size = 1000

# Execute Training

for epoch in range(n_epoch):
    sum_loss = 0
    perm = np.random.permutation(60000)
    for i in range(0, 60000, batch_size):
        x = Variable(x_train[perm[i:i+batch_size]])
        t = Variable(t_train[perm[i:i+batch_size]])
        y = model.forward(x)
        model.cleargrads()
        loss = F.softmax_cross_entropy(y, t)
        loss.backward()
        optimizer.update()
        sum_loss += loss.data*batch_size

    print("epoch: {}, mean loss: {}".format(epoch, sum_loss/60000))

# Execute Test
cnt = 0
for i in range(10000):
    x = Variable(np.array([x_test[i]], dtype=np.float32))
    t = t_test[i]
    y = model.forward(x)
    y = np.argmax(y.data[0])
    if t == y:
        cnt += 1

# Display Result
print("accuracy: {}".format(cnt/(10000)))

このコード、中間層のノード数やエポック数を変更している以外はほぼ下記のAI研究所さんブログのコピペです。
こちらのブログにはChainerを始めたばかりの頃に結構お世話になりました。
ai-kenkyujo.com

このコードをベースにして順番にPyTorch用に書き換えて行きます。

  1. import文の変更


    ChainerPyTorch

    # Numpy
    import numpy as np
    # Chainer
    import chainer
    import chainer.links as L
    import chainer.functions as F
    from chainer import Chain,optimizers,Variable

    # Numpy
    import numpy as np
    # Chainer
    # ★Chainerのdatasetsからmnistを引っ張るためchainerは残す
    import chainer
    # ★Pytorchをインポート
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import torch.optim as optim


    ここでimport chainerを残しているのは、chainer.datasetsから取得したMNISTデータを使ってPyTorchで学習させるためです。PyTorchにもMNISTのデータセットを取得する機能があるのですが、今回はchainer版と同じ形式のデータをPyTorchで学習させるサンプルのため、chainer.datasetsからのデータセット取得をそのまま残します。 chainerの場合、データをnumpyの行列にした後Variables()で包んでニューラルネットワークのモデルに渡しますが、このnumpyの行列にするところまでは同じ流れで進めます。 あくまでもchainer用に構築した学習環境・データをモデルのみPyTorchで定義したモデルに差し替えて動かすための基本的な動作確認をすることが、この移植試行の目的です。
  2. numpy行列化したMNISTデータの準備


    ChainerPyTorch

    train, test = chainer.datasets.get_mnist()
    
    x_train = []
    x_test  = []
    t_train = []
    t_test  = []
    
    for line in train:
        x_train.append(line[0])
        t_train.append(int(line[-1]))
    
    for line in test:
        x_test.append(line[0])
        t_test.append(int(line[-1]))
    
    x_train = np.array(x_train,dtype=np.float32)
    x_test = np.array(x_test,dtype=np.float32)
    t_train = np.array(t_train,dtype=np.int32)
    t_test = np.array(t_test,dtype=np.int32)

    train, test = chainer.datasets.get_mnist()
    
    x_train = []
    x_test  = []
    t_train = []
    t_test  = []
    
    for line in train:
        x_train.append(line[0])
        t_train.append(int(line[-1]))
    
    for line in test:
        x_test.append(line[0])
        t_test.append(int(line[-1]))
    
    x_train = np.array(x_train,dtype=np.float32)
    x_test = np.array(x_test,dtype=np.float32)
    # ★PyTorchでは分類クラスはintではなくlong
    t_train = np.array(t_train,dtype=np.int64)
    t_test = np.array(t_test,dtype=np.int64)


    ほぼ同じなのですが、教師データ(t_train及びt_test: MNISTの0〜9までのラベル)はnp.int32だとPyTorchで学習時にエラーが出るためnp.int64にしています。PyTorchでは分類データのラベルはintではなくlongにする必要がある様です。これは一度np.int32のまま実行してみたらエラーが吐き出されました。
  3. モデル定義


    ChainerPyTorch

    # DNN class
    class DNN(Chain):
        def __init__(self):
            super(DNN, self).__init__(
                l1 = L.Linear(None,1000),
                l2 = L.Linear(None,1000),
                l3 = L.Linear(None,10)
            )
        def forward(self,x):
            h1 = F.relu(self.l1(x))
            h2 = F.relu(self.l2(h1))
            h3 = self.l3(h2)
            return h3

    # DNN class
    class DNN(nn.Module):
        def __init__(self):
            super(DNN,self).__init__()
            self.l1 = nn.Linear(784,1000)
            self.l2 = nn.Linear(1000,1000)
            self.l3 = nn.Linear(1000,10)
        def forward(self,x):
            h1 = F.relu(self.l1(x))
            h2 = F.relu(self.l2(h1))
            h3 = self.l3(h2)
            return h3


    こちらもほぼ同じなのですが、ChainerではLinear()で全結合を作る際に前の段のネットワークから受けるノード数はNoneを指定しておくと自動的に前の段の出力に合わせてくれるのですが、PyTorchは数値で設定が必要です。
  4. モデル生成、オプティマイザ−セットアップ


    ChainerPyTorch

    # Create DNN class instance
    model = DNN()
    
    # Set optimizer
    optimizer = optimizers.Adam()
    optimizer.setup(model)

    # Create DNN class instance
    model = DNN()
    
    # Set optimizer
    # ★optimizerをセット
    optimizer = optim.Adam(model.parameters())
    


    モデル生成はモデルのクラスインスタンスを作るだけなので同じです。 オプティマイザ−へのモデルの設定が違う書き方になります。
  5. 学習ループ


    ChainerPyTorch

    # Number of epochs
    n_epoch = 10
    
    # batch size
    batch_size = 1000
    
    # Execute Training
    
    for epoch in range(n_epoch):
        sum_loss = 0
        perm = np.random.permutation(60000)
        for i in range(0, 60000, batch_size):
            x = Variable(x_train[perm[i:i+batch_size]])
            t = Variable(t_train[perm[i:i+batch_size]])
            y = model.forward(x)
            model.cleargrads()
            loss = F.softmax_cross_entropy(y, t)
            loss.backward()
            optimizer.update()
            sum_loss += loss.data*batch_size
    
        print("epoch: {}, mean loss: {}".format(epoch, sum_loss/60000))
    

    # Number of epochs
    n_epoch = 10
    
    # batch size
    batch_size = 1000
    
    # Execute Training
    
    for epoch in range(n_epoch):
        sum_loss = 0
        perm = np.random.permutation(60000)
        for i in range(0, 60000, batch_size):
    # ★x_trainとy_trainをpytorchのTensorに変換
    #   ここで変換は効率悪い気もするけど..元のコードベースを崩さないため
            x = torch.from_numpy(x_train[perm[i:i+batch_size]])
            t = torch.from_numpy(t_train[perm[i:i+batch_size]])
            y = model.forward(x)
    # ★cleargrads()はzero_grad()に
            model.zero_grad()
    # ★softmax_cross_entropy()→cross_entropy()
            loss = F.cross_entropy(y,t)
            loss.backward()
    #  ★optimizer.update()はoptimizer.step()に
            optimizer.step()
            sum_loss += loss.data*batch_size
    
        print("epoch: {}, mean loss: {}".format(epoch, sum_loss/60000))
    

    学習データと教師データをVariablesで包んでいるところをtorch.from_numpy()でPyTorchのTensor化します。 後は損失関数のsoftmax_cross_entropy()の関数名が変わるとか、optimizer.update()をoptimizer.setup()に変える程度です。 これで学習が出来ます。 なお自分はChainerのTrainerは使わない(使えないが正しい...)ので、学習ループはいつも手で組んでいます。
  6. 学習結果の検証(テスト実行)


    ChainerPyTorch

    # Execute Test
    cnt = 0
    for i in range(10000):
        x = Variable(np.array([x_test[i]], dtype=np.float32))
        t = t_test[i]
        y = model.forward(x)
        y = np.argmax(y.data[0])
        if t == y:
            cnt += 1
    
    # Display Result
    print("accuracy: {}".format(cnt/(10000)))
    

    # Execute Test
    cnt = 0
    for i in range(10000):
    # ★numpy to Tensor
        x = torch.from_numpy(np.array([x_test[i]], dtype=np.float32))
        t = t_test[i]
        y = model.forward(x)
    # ★Tensorからnumpy()に戻す,めんどくさい...
        y = y.detach().numpy()
        y = np.argmax(y[0])
        if t == y:
            cnt += 1
    
    # Display Result
    print("accuracy: {}".format(cnt/(10000)))


    推論結果をnumpyのargmax()に渡して判定するところで、Chainerだとy.data[0]を直接numpy.argmax()に渡せますが、PyTorchでは一旦明示的にnumpy()に変換してから渡す必要がありました。
  7. 最後に動かしてみる

    Chainer

    (py36) toy@DELL-0001:~/chainer_test/mnist$ python mnist_test.py
    epoch: 0, mean loss: 0.4059475188453992
    epoch: 1, mean loss: 0.13042979886134465
    epoch: 2, mean loss: 0.07960983558247486
    epoch: 3, mean loss: 0.053931693515429896
    epoch: 4, mean loss: 0.038755741560210784
    epoch: 5, mean loss: 0.02595699280500412
    epoch: 6, mean loss: 0.01904308699692289
    epoch: 7, mean loss: 0.013304473071669539
    epoch: 8, mean loss: 0.010054900473915041
    epoch: 9, mean loss: 0.007003770768642425
    accuracy: 0.9808



    PyTorch

    (py36) toy@DELL-0001:~/chainer_test/mnist$ python mnist_test_torch.py
    epoch: 0, mean loss: 0.49522629380226135
    epoch: 1, mean loss: 0.17571552097797394
    epoch: 2, mean loss: 0.11051174998283386
    epoch: 3, mean loss: 0.07754945755004883
    epoch: 4, mean loss: 0.05589161440730095
    epoch: 5, mean loss: 0.04250581935048103
    epoch: 6, mean loss: 0.03223542869091034
    epoch: 7, mean loss: 0.024668868631124496
    epoch: 8, mean loss: 0.01830950565636158
    epoch: 9, mean loss: 0.014577378518879414
    accuracy: 0.9816
    



    オプティマイザ−のデフォルトのパラメータが異なるのか、損失の減り方が異なりますが、10エポックでほぼ同等の精度になっています。
    ChainerとPyTorchの使い方は非常に似ていると言われており、モデルの定義や学習ループ等はほぼおなじ構造のまま移植出来るのですが、やはり細かい関数名の違いやVariableとTensorの違い等があり、手作業での移植はなかなか面倒くさいです。
    Chainer開発元のPreferred Networksで、ChainerのパラメータをPyTorchのパラメータに変換したりChainerのモデルをPyTorchのモデルに変換するMigration Toolを公開しているので、時間があれば、こちらも試してみたいと思います。
    https://chainer.github.io/migration-guide/
    GitHub - chainer/chainer-pytorch-migration: Chainer/PyTorch Migration Library


後、試したい事としては、

くらいかな。

最後にPyTorch版MNISTソース全体を載せておきます。

# -*- Coding: utf-8 -*-

# Numpy
import numpy as np
# Chainer
# ★Chainerのdatasetsからmnistを引っ張るためchainerは残す
import chainer
# ★Pytorchをインポート
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

train, test = chainer.datasets.get_mnist()

x_train = []
x_test  = []
t_train = []
t_test  = []

for line in train:
    x_train.append(line[0])
    t_train.append(int(line[-1]))

for line in test:
    x_test.append(line[0])
    t_test.append(int(line[-1]))

x_train = np.array(x_train,dtype=np.float32)
x_test = np.array(x_test,dtype=np.float32)
# ★PyTorchでは分類クラスはintではなくlong
t_train = np.array(t_train,dtype=np.int64)
t_test = np.array(t_test,dtype=np.int64)


# DNN class
class DNN(nn.Module):
    def __init__(self):
        super(DNN,self).__init__()
        self.l1 = nn.Linear(784,1000)
        self.l2 = nn.Linear(1000,1000)
        self.l3 = nn.Linear(1000,10)
    def forward(self,x):
        h1 = F.relu(self.l1(x))
        h2 = F.relu(self.l2(h1))
        h3 = self.l3(h2)
        return h3

# Create DNN class instance
model = DNN()

# Set optimizer
# ★optimizerをセット
optimizer = optim.Adam(model.parameters())

# Number of epochs
n_epoch = 10

# batch size
batch_size = 1000

# Execute Training

for epoch in range(n_epoch):
    sum_loss = 0
    perm = np.random.permutation(60000)
    for i in range(0, 60000, batch_size):
# ★x_trainとy_trainをpytorchのTensorに変換
#   ここで変換は効率悪い気もするけど..元のコードベースを崩さないため
        x = torch.from_numpy(x_train[perm[i:i+batch_size]])
        t = torch.from_numpy(t_train[perm[i:i+batch_size]])
        y = model.forward(x)
# ★cleargrads()はzero_grad()に
        model.zero_grad()
# ★softmax_cross_entropy()→cross_entropy()
        loss = F.cross_entropy(y,t)
        loss.backward()
#  ★optimizer.update()はoptimizer.step()に
        optimizer.step()
        sum_loss += loss.data*batch_size

    print("epoch: {}, mean loss: {}".format(epoch, sum_loss/60000))

# Execute Test
cnt = 0
for i in range(10000):
# ★numpy to Tensor
    x = torch.from_numpy(np.array([x_test[i]], dtype=np.float32))
    t = t_test[i]
    y = model.forward(x)
# ★Tensorからnumpy()に戻す,めんどくさい...
    y = y.detach().numpy()
    y = np.argmax(y[0])
    if t == y:
        cnt += 1

# Display Result
print("accuracy: {}".format(cnt/(10000)))

参考にしたサイト:
qiita.com
qiita.com

pytorch.org