Kerasからpytorchへ移植

少し前に書いたKerasでのhello world的なコードをpytorchへ移植してみたときのメモ。

MNISTの文字認識のあれです。

準備
コード

準備

インストールしたものは、torch / torchvision / pillow / tqdm / matplotlibの5つです。全部pip installからインストールできました。

コード

いろいろネットをさまよいながらコピペしながら移植したので、いろんなコードの張り合わせになってます。やはりKerasよりコードは長くなっちゃいますね。

基本的に用語が同じなので、まぁなんとなくで移植はできした。学習時間が長いので、学習からその学習データの保存までと、読み込みから推論するまでの構成に分けました。

from tqdm import tqdm
import matplotlib.pyplot as plt
from PIL import Image
 
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
 
def train_model(model, train_loader, criterion, optimizer, device='cpu'):
 
    train_loss = 0.0
    num_train = 0
    correct = 0
 
    # 学習モデルに変換
    model.train()
 
    # 進捗の表示
    pbar = tqdm(train_loader)
     
    for i, (images, labels) in enumerate(pbar):
        # batch数をカウント
        num_train += len(labels)
 
        images, labels = images.to(device), labels.to(device)
 
        # 勾配を初期化
        optimizer.zero_grad()
 
        # 推論(順伝播)
        outputs = model(images)
         
        # 推論の結果を取得
        preds = outputs.argmax(dim=1)
 
        # 損失の算出
        loss = criterion(outputs, labels)
 
        # 誤差逆伝播
        loss.backward()
 
        # パラメータの更新
        optimizer.step()
 
        # lossを加算
        train_loss += loss.item()
 
        # accuracyを加算
        correct += (preds == labels).sum()
     
    # lossの平均値を取る
    train_loss = train_loss / num_train
 
    # 正解率を取る
    accuracy = correct / num_train
 
    return train_loss, accuracy
 
 
def test_model(model, test_loader, criterion, device='cpu'):
 
    test_loss = 0.0
    num_test = 0
    correct = 0
 
    # modelを評価モードに変更
    model.eval()
 
    with torch.no_grad(): # 勾配計算の無効化
        for i, (images, labels) in enumerate(test_loader):
            num_test += len(labels)
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
             
            preds = outputs.argmax(dim=1)
            correct += (preds == labels).sum()
             
            loss = criterion(outputs, labels)
            test_loss += loss.item()
         
        test_loss = test_loss / num_test
        accuracy = correct / num_test
 
    return test_loss, accuracy
 
 
def lerning(model, train_loader, test_loader, criterion, opimizer, num_epochs, device='cpu'):
 
    train_loss_list = []
    test_loss_list = []
 
    # epoch数分繰り返す
    for epoch in range(1, num_epochs+1, 1):
 
        train_loss, train_accuracy = train_model(model, train_loader, criterion, optimizer, device=device)
        test_loss, test_accuracy = test_model(model, test_loader, criterion, device=device)
         
        print("epoch : {}, train_accuracy : {:.5f}, test_accuracy : {:.5f}" .format(epoch, train_accuracy, test_accuracy))
        print("epoch : {}, train_loss : {:.5f}, test_loss : {:.5f}" .format(epoch, train_loss, test_loss))
 
        train_loss_list.append(train_loss)
        test_loss_list.append(test_loss)
     
    return train_loss_list, test_loss_list
 
 
 
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
         
        self.layer1 = nn.Sequential(
                nn.Conv2d(1, 20, 5, padding="same"),
                nn.ReLU(),
                nn.MaxPool2d(2),
                nn.Dropout(0.25)
            )
         
        self.layer2 = nn.Sequential(
                nn.Conv2d(20, 50, 5, padding="same"),
                nn.ReLU(),
                nn.MaxPool2d(2),
                nn.Dropout(0.25)
            )
         
        self.flaten1 = nn.Flatten()
         
        self.layer3 = nn.Sequential(
                nn.Linear(7*7*50, 256),
                nn.Sigmoid(),
                nn.Linear(256, 128),
                nn.Sigmoid(),
                nn.Dropout(0.2),
                nn.Linear(128, 10),
                nn.Softmax()
            )
 
 
    def forward(self, x): # x : 入力
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.flaten1(x)
        x = self.layer3(x)
        return x
 
 
 
#訓練データ
train_dataset = torchvision.datasets.MNIST(root='./data',
                                           train=True,
                                           transform=transforms.ToTensor(),
                                           download = True)
#検証データ
test_dataset = torchvision.datasets.MNIST(root='./data',
                                           train=False,
                                           transform=transforms.ToTensor(),
                                           download = True)
 
batch_size = 128
 
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
 
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
 
model = Net()
print(model)
# print(list(model.parameters()))
 
# 学習に使うパラメータ
optimizer = optim.Adam(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
num_epochs = 5
 
train_loss_list, test_loss_list = lerning(model, train_loader, test_loader, criterion, optimizer, num_epochs)
 
# 学習データの保存
torch.save(model.state_dict(), "./mnist.pth")

LossとかAccuracyも自分で計算しないといけないんですね。TorchEvalっていう別ライブラリはあるみたいですが。

これを実行すると

100%|██████████| 469/469 [00:28<00:00, 16.41it/s]
epoch : 1, train_accuracy : 0.85400, test_accuracy : 0.96820
epoch : 1, train_loss : 0.01261, test_loss : 0.01181
100%|██████████| 469/469 [00:32<00:00, 14.58it/s]
epoch : 2, train_accuracy : 0.95360, test_accuracy : 0.96230
epoch : 2, train_loss : 0.01180, test_loss : 0.01185
100%|██████████| 469/469 [00:30<00:00, 15.47it/s]
epoch : 3, train_accuracy : 0.94738, test_accuracy : 0.96850
epoch : 3, train_loss : 0.01185, test_loss : 0.01180
100%|██████████| 469/469 [00:32<00:00, 14.39it/s]
epoch : 4, train_accuracy : 0.94225, test_accuracy : 0.96400
epoch : 4, train_loss : 0.01188, test_loss : 0.01184
100%|██████████| 469/469 [00:38<00:00, 12.21it/s]
epoch : 5, train_accuracy : 0.94200, test_accuracy : 0.95290
epoch : 5, train_loss : 0.01189, test_loss : 0.01193

なんとなくうまくいっているように見えます。1epoch目から全然伸びない傾向も似てます。ただ正答率がKerasで作ってた時よりちょっと低いなぁ。何かミスってるかな。

続いて、訓練データの先頭６つからの推論

# 学習データの読み込み
model.load_state_dict(torch.load("./mnist.pth", map_location=torch.device('cpu')))
 
# test data 6枚だけ画像の表示と推論
fig = plt.figure()
for i in range(6):
    image, label = test_dataset[i]
    image = image.view(1, 1, 28, 28).to('cpu')
 
    # 推論
    prediction_label = torch.argmax(model(image))
    print('label : {}\n Prediction : {}'.format(label, prediction_label))
 
    fig.add_subplot(2,3,i+1)
    plt.imshow(image.detach().to('cpu').numpy().reshape(28, 28), cmap = "gray")

なんかはまったのが、image.view(1,1,28,28)の記述。４次元のテンソルにしておかないとflattenのところで引っかかりました。訓練時には１次元目にbatch数が入るみたいなのですが、推論時にはそれが消えていて掛け算のサイズが合わないよ。って怒られました。

自分で手書きした外にある画像ファイル（２８x２８のグレースケールのTiff）からの推論

# 画像データの読み込みとテンソル化
to_tensor = transforms.ToTensor()
testimage = to_tensor(Image.open("testdata.tif")).unsqueeze(0)
 
# 推論
predict = torch.argmax(model(testimage))
print("predict: {}".format(predict))
plt.imshow(testimage.detach().to('cpu').numpy().reshape(28, 28), cmap = "gray")