PytorchでCNN(画像分類)を書いてみる。

PytorchでCNN(画像分類)を書いてみる。

以下を参考にCIFAR10 datasetを用いたCNNを書いてみる。
pytorch.org

ソースコード

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchsummary import summary
import matplotlib.pyplot as plt
import numpy as np

# 学習回数
EPOCH = 4

PATH = 'データセットを保存するディレクトリ'

# 前処理の定義と正規化
transform = transforms.Compose(
        [transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])

# トレーニング用データセットのダウンロードと前処理
# 引数
# root : データセットを保存するディレクトリ。
# train : トレーニング用のデータを取得するかどうかを選択。
# download:参照したディレクトリにDatasetがない場合ダウンロードするかどうかを決める。
# transform:定義した前処理を渡す。これにより、Dataset内のdataを参照する際にその前処理を自動で行ってくれる。
trainset = torchvision.datasets.CIFAR10(root=PATH, train=True,download=True, transform=transform)

# トレーニング用のDataloaderの作成
# 引数
# 第1引数は先程取得したDatasetを入れる
# batch_size：1回のtrainingまたはtest時にまとめて何個のdataを使用するかを選択。(datasetの全data数を割り切れる値にしなければならない)
# shuffle：dataの参照の仕方をランダムにする。
# num_workers:並列処理を行う。2以上の場合はその数だけ並列処理を行う。
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

# テスト用のデータセットのダウンロードと前処理
testset = torchvision.datasets.CIFAR10(root=PATH, train=False, download=True, transform=transform)

# テスト用のDataloaderの作成
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)
      
### ネットワーク(CNN)の定義 ###
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5) # nn.Conv2d(input, output, kernel_size)
        self.pool  = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1   = nn.Linear(16*5*5, 120) # nn.Linear(input, output)
        self.fc2   = nn.Linear(120, 84)
        self.fc3   = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# ネットワークのインスタンス化
net = Net()

# loss関数の定義
criterion = nn.CrossEntropyLoss()

# 最適化関数の定義
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

### トレーニング ###
for epoch in range(EPOCH):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # データ[input, labels]を取得
        inputs, labels = data

        # 均配の初期化
        optimizer.zero_grad()

        #フォワード処理forward関数の実行）
        outputs = net(inputs)

        # lossを計算
        loss = criterion(outputs, labels)

        # lossから均配を計算
        loss.backward()

        # 最適化(均配からパラメータの更新)
        optimizer.step()

        # 学習過程を表示
        running_loss += loss.item()
        if i % 2000 == 1999: # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' % 
                    (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

# 次回から学習しなくていいように、学習済みモデルのパラメータを"net_00.prmとして保存"
params = net.state_dict()
torch.save(params, "net_00.prm", pickle_protocol=4)

print('Finished Training')

### 学習済みモデルのテスト ###
correct = 0
total = 0

with torch.no_grad(): # パラメータ保持の停止（test時には必要ないので。）
    for data in testloader:
        images, labels = data

        # 各ラベルの確率がoutputに出力される
        outputs = net(images)

        # torch.maxは最大値とそのインデックス(ラベル)を返す
        # 最大値は無視してラベルのみを保存
        _, predicted = torch.max(outputs.data, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()

# 正解率を出力
print('Accuracy of the network on the 10000 test images: %d %%' %
        (100 * correct / total))

実行結果

Accuracy of the network on the 10000 test images: 56 %