Kerasからpytorchへ移植

python

少し前に書いた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で作ってた時よりちょっと低いなぁ。何かミスってるかな。

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

# 学習データの読み込み
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)の記述。4次元のテンソルにしておかないとflattenのところで引っかかりました。訓練時には1次元目にbatch数が入るみたいなのですが、推論時にはそれが消えていて掛け算のサイズが合わないよ。って怒られました。

自分で手書きした外にある画像ファイル(28x28のグレースケールの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")

なのでここでもunsqueezeで次元追加。

まぁ忘れないようにメモしておきます。

python画像処理
スポンサーリンク
キャンプ工学

コメント

タイトルとURLをコピーしました