少し前に書いた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")
なのでここでもunsqueezeで次元追加。
まぁ忘れないようにメモしておきます。
コメント