【PyTorch始めました】Chainerで書いた畳み込みニューラルネットワークをPyTorchに移植する

[technology]畳み込みニューラルネットワークのChainer to PyTorch対訳

MNISTの手描き文字認識のChainer-PyTorchの対訳の畳み込みニューラルネットワーク版を作りました。
まず、Chainer版の元ソースとそれをPyTorch向けに書き換えたソースをそのまま載せます。

前回と同じく、MNISTのデータはどちらの場合もchainer.datasetsから取得します。
取得したMNISTデータをnumpyの行列化した後、更にPyTorchのTensor化する流れは同じですが、畳み込みニューラルネットワークへのデータ流し込みのため、MNISTの各画像データを1(チャンネル) x 28 x 28の形に形状変更(reshape)をしています。

後はChainerで書かれた畳み込みニューラルネットワークのモデルをPyTorchのモデルに置き換えます。

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
from chainer import serializers
import matplotlib.pyplot as plt

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)

x_train = x_train.reshape(len(x_train),1,28,28)
x_test = x_test.reshape(len(t_test),1,28,28)

# DNN class
class DNN(Chain):
    def __init__(self):
        super(DNN, self).__init__(
            cv1 = L.Convolution2D(1,5,4),
            cv2 = L.Convolution2D(5,10,4),
            l1 = L.Linear(None,100),
            l2 = L.Linear(None,10)
        )
    def forward(self,x):
        c = F.max_pooling_2d(F.relu(self.cv1(x)),2)
        c = F.max_pooling_2d(F.relu(self.cv2(c)),2)
        h = F.relu(self.l1(c))
        h = self.l2(h)
        return h

# Create DNN class instance
model = DNN()

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

# Number of epochs
n_epoch = 20

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

PyTorch移植ソース

# -*- 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)

x_train = x_train.reshape(len(x_train),1,28,28)
x_test = x_test.reshape(len(t_test),1,28,28)

        
# DNN class
class DNN(nn.Module):
    def __init__(self):
        super(DNN,self).__init__()
        self.cv1 = nn.Conv2d(1,5,4)
        self.cv2 = nn.Conv2d(5,10,4)
        self.l1 = nn.Linear(10 * 4 * 4,100)
        self.l2 = nn.Linear(100,10)
    def forward(self,x):
        c = F.max_pool2d(F.relu(self.cv1(x)),2)
        c = F.max_pool2d(F.relu(self.cv2(c)),2)
        # Chainerと異なり、全結合層に入力する前にshape変換が必要
        c = c.view(-1,10 * 4 * 4)
        h = F.relu(self.l1(c))
        h = self.l2(h)
        return h

# Create DNN class instance
model = DNN()

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

# Number of epochs
n_epoch = 20

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

モデル部分を比較すると以下の通りです。

Chainer	PyTorch
# DNN class class DNN(Chain): def __init__(self): super(DNN, self).__init__( cv1 = L.Convolution2D(1,5,4), cv2 = L.Convolution2D(5,10,4), l1 = L.Linear(None,100), l2 = L.Linear(None,10) ) def forward(self,x): c = F.max_pooling_2d(F.relu(self.cv1(x)),2) c = F.max_pooling_2d(F.relu(self.cv2(c)),2) h = F.relu(self.l1(c)) h = self.l2(h) return h	# DNN class class DNN(nn.Module): def __init__(self): super(DNN,self).__init__() self.cv1 = nn.Conv2d(1,5,4) self.cv2 = nn.Conv2d(5,10,4) self.l1 = nn.Linear(10 * 4 * 4,100) self.l2 = nn.Linear(100,10) def forward(self,x): c = F.max_pool2d(F.relu(self.cv1(x)),2) c = F.max_pool2d(F.relu(self.cv2(c)),2) # Chainerと異なり、全結合層に入力する前にshape変換が必要 c = c.view(-1,10 * 4 * 4) h = F.relu(self.l1(c)) h = self.l2(h) return h

ほぼ同じ様な定義なのですが、Chainerの全結合層はLinear(None,100)等と書いておくと全結合層の入力側のノード数は前段のノード数に自動的に合わせこんでくれます。また前段の出力が何次元であっても、1次元に自動展開してくれます。この便利な機能のおかげで、畳み込みとMAX Poolingを繰り返した結果として入力のノード数が何段になるかを明示的に計算して定義しなくても良いので非常に楽です。
PyTorchの場合、

• 全結合層に入力する前に畳み込み層の出力を1次元に明示的に形状変換が必要(view()を使う)
• 畳み込み層の出力サイズを計算して全結合層の入力ノード数として明示的に定義が必要

の2点がChainerと異なります。

畳み込み層の出力サイズの計算は例えば
deepage.net
等に解説されており、計算することが出来ますが、自分の場合はpythonをREPLモードで立ち上げて、実際に畳み込み演算を実行してサイズをチェックして、モデル定義に反映しています。以下は今回の畳み込みニューラルネットワークの出力サイズのPyTorchを使った確認例です。

>>> import numpy as np
>>> import torch
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> conv1 = nn.Conv2d(1, 5, 4)
>>> conv2 = nn.Conv2d(5,10,4)
>>> x = np.random.rand(28 * 28).astype(np.float32)
>>> x = x.reshape(1,1,28,28)
>>> x = torch.from_numpy(x)
>>> x1 = conv1(x)
>>> x2 = F.max_pool2d(x1,2)
>>> x3 = conv2(x2)
>>> x4 = F.max_pool2d(x3,2)
>>> x4.shape
torch.Size([1, 10, 4, 4])

こんな感じで、10*4*4=160になることが判ります。

ちなみにChainerではこれをやる時に x = torch.from_numpy(x)は不要で、numpyの行列をそのまま畳み込み等の関数に食わせる事ができます。
で、今更気づいたのですが、今まで自分はChainerで訓練したネットワークを使って推論する際も、numpy(あるいはcupy)の行列をVariablesで包んでいたのですが、推論の場合は微分が不要なのでVariablesは要らないのですね。意外とちゃんと理解できていないなあ。

実際にChainer版とPyTorch版でMNISTを学習させた例です。全結合版と同じく若干最適化の進み方が異なりますが、今回のモデルの場合は、20epochでほぼ同等の損失と正解率になりました。

第７世代のCore i7では同じ20epochの学習ではPyTorchの方が圧倒的に速いです。PyTorch 1分38秒に対してChainerは5分48秒かかりました。
次はGPUでPyTorch版の畳み込みニューラルネットワークを動かしてみたいと思います。
Chainer版の実行例:

$ python mnist_cnn_test.py
epoch: 0, mean loss: 1.1072022398312886
epoch: 1, mean loss: 0.2801803067326546
epoch: 2, mean loss: 0.1904003771642844
epoch: 3, mean loss: 0.14446297784646353
epoch: 4, mean loss: 0.11799863787988822
epoch: 5, mean loss: 0.10124019905924797
epoch: 6, mean loss: 0.0898616046955188
epoch: 7, mean loss: 0.07999170552939176
epoch: 8, mean loss: 0.07271959017962218
epoch: 9, mean loss: 0.06986213885247708
epoch: 10, mean loss: 0.06315046300490697
epoch: 11, mean loss: 0.05993131225307782
epoch: 12, mean loss: 0.055548572726547715
epoch: 13, mean loss: 0.055429142465194065
epoch: 14, mean loss: 0.04930390634884437
epoch: 15, mean loss: 0.047334119336058696
epoch: 16, mean loss: 0.04560358158002297
epoch: 17, mean loss: 0.04443015136445562
epoch: 18, mean loss: 0.041684689931571484
epoch: 19, mean loss: 0.039714485468963785
accuracy: 0.986

PyTorch版の実行例:

$ python mnist_cnn_test_torch.py
epoch: 0, mean loss: 1.6421406269073486
epoch: 1, mean loss: 0.4664209187030792
epoch: 2, mean loss: 0.32162800431251526
epoch: 3, mean loss: 0.251921683549881
epoch: 4, mean loss: 0.2077903300523758
epoch: 5, mean loss: 0.17509377002716064
epoch: 6, mean loss: 0.1511804163455963
epoch: 7, mean loss: 0.133607879281044
epoch: 8, mean loss: 0.11779090017080307
epoch: 9, mean loss: 0.10718908905982971
epoch: 10, mean loss: 0.09853501617908478
epoch: 11, mean loss: 0.09329884499311447
epoch: 12, mean loss: 0.08585326373577118
epoch: 13, mean loss: 0.08047447353601456
epoch: 14, mean loss: 0.07560740411281586
epoch: 15, mean loss: 0.07071471214294434
epoch: 16, mean loss: 0.0669841319322586
epoch: 17, mean loss: 0.06552083790302277
epoch: 18, mean loss: 0.061434123665094376
epoch: 19, mean loss: 0.060454342514276505
accuracy: 0.9824

【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文の変更
   
   

   Chainer	PyTorch
   # 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データの準備
   
   

   Chainer	PyTorch
   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. モデル定義
   
   

   Chainer	PyTorch
   # 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. モデル生成、オプティマイザ−セットアップ
   
   

   Chainer	PyTorch
   # 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. 学習ループ
   
   
   

   Chainer	PyTorch
   # 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. 学習結果の検証(テスト実行)
   
   

   Chainer	PyTorch
   # 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
   
   
   

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

• モデルを全結合→畳み込みニューラルネットワークに変更してみる
  
• CPUではなく、GPUで回してみる

くらいかな。

最後に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