【Pytorch】CNNでmnistクラスの分類を学習するコード 【mnist】
mnistを用いてクラス分類を行うモデルをPytorchを使って学習するためのコードをまとめておく。
モデルの各レイヤーの定義をしているコードについては、以下のページにまとめてあるのでまずはそこを参照して実装する。
モデルの構造を定義
次に、以下でモデルの構造を定義する。
バッチサイズやエポック数、最適化手法も定義している。
conv_net = nn.Sequential(
Conv((20, 1, 5, 5), F.relu), # 28x28x 1 -> 24x24x20
Pooling((2, 2)), # 24x24x20 -> 12x12x20
Conv((50, 20, 5, 5), F.relu), # 12x12x20 -> 8x 8x50
Pooling((2, 2)), # 8x 8x50 -> 4x 4x50
Flatten(),
Dense(4*4*50, 10)
)
batch_size = 100
n_epochs = 10
lr = 0.02
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
conv_net.to(device)
optimizer = optim.SGD(conv_net.parameters(), lr=lr)mnistのデータを読み込む
以下でmnistのデータを読み込んで学習に使えるようにする。
# torchvisionのdatasetsを使ってMNISTのデータを取得
# ミニバッチ化や前処理などの処理を行ってくれるDataLoaderを定義
dataloader_train = torch.utils.data.DataLoader(
datasets.MNIST('./data/mnist', train=True, download=True, transform=transforms.ToTensor()),
batch_size=batch_size,
shuffle=True
)
dataloader_valid = torch.utils.data.DataLoader(
datasets.MNIST('./data/mnist', train=False, download=True, transform=transforms.ToTensor()),
batch_size=batch_size,
shuffle=False
)モデルを学習する
最後に、以下のプログラムを実行してモデルの学習を行う。
for epoch in range(n_epochs):
losses_train = []
losses_valid = []
conv_net.train()
n_train = 0
acc_train = 0
for x, t in dataloader_train:
n_train += t.size()[0]
conv_net.zero_grad() # 勾配の初期化
x = x.to(device) # テンソルをGPUに移動
t_hot = torch.eye(10)[t] # 正解ラベルをone-hot vector化
t_hot = t_hot.to(device) # 正解ラベルとone-hot vectorをそれぞれGPUに移動
y = conv_net.forward(x) # 順伝播
loss = -(t_hot*torch.log_softmax(y, dim=-1)).sum(axis=1).mean() # 誤差(クロスエントロピー誤差関数)の計算
loss.backward() # 誤差の逆伝播
optimizer.step() # パラメータの更新
pred = y.argmax(1) # 最大値を取るラベルを予測ラベルとする
acc_train += (pred.to("cpu") == t).float().sum().item()
losses_train.append(loss.tolist())
conv_net.eval()
n_val = 0
acc_val = 0
for x, t in dataloader_valid:
n_val += t.size()[0]
x = x.to(device) # テンソルをGPUに移動
t_hot = torch.eye(10)[t] # 正解ラベルをone-hot vector化
t_hot = t_hot.to(device) # 正解ラベルとone-hot vectorをそれぞれGPUに移動
y = conv_net.forward(x) # 順伝播
loss = -(t_hot*torch.log_softmax(y, dim=-1)).sum(axis=1).mean() # 誤差(クロスエントロピー誤差関数)の計算
pred = y.argmax(1) # 最大値を取るラベルを予測ラベルとする
acc_val += (pred.to("cpu") == t).float().sum().item()
losses_valid.append(loss.tolist())
print('EPOCH: {}, Train [Loss: {:.3f}, Accuracy: {:.3f}], Valid [Loss: {:.3f}, Accuracy: {:.3f}]'.format(
epoch,
np.mean(losses_train),
acc_train/n_train,
np.mean(losses_valid),
acc_val/n_val
))
# 出力
# EPOCH: 0, Train [Loss: 0.471, Accuracy: 0.866], Valid [Loss: 0.244, Accuracy: 0.928]
# EPOCH: 1, Train [Loss: 0.225, Accuracy: 0.935], Valid [Loss: 0.168, Accuracy: 0.952]
# EPOCH: 2, Train [Loss: 0.167, Accuracy: 0.953], Valid [Loss: 0.130, Accuracy: 0.963]
# EPOCH: 3, Train [Loss: 0.135, Accuracy: 0.961], Valid [Loss: 0.112, Accuracy: 0.968]
# EPOCH: 4, Train [Loss: 0.116, Accuracy: 0.967], Valid [Loss: 0.102, Accuracy: 0.968]
# EPOCH: 5, Train [Loss: 0.101, Accuracy: 0.971], Valid [Loss: 0.080, Accuracy: 0.976]
# EPOCH: 6, Train [Loss: 0.092, Accuracy: 0.974], Valid [Loss: 0.076, Accuracy: 0.976]
# EPOCH: 7, Train [Loss: 0.084, Accuracy: 0.975], Valid [Loss: 0.072, Accuracy: 0.978]
# EPOCH: 8, Train [Loss: 0.078, Accuracy: 0.978], Valid [Loss: 0.061, Accuracy: 0.982]
# EPOCH: 9, Train [Loss: 0.074, Accuracy: 0.979], Valid [Loss: 0.060, Accuracy: 0.981]