データセットの準備

import torchvision

training_dataset = torchvision.datasets.CIFAR10("./cifar10_data", train=True, download=True,
                                           transform=torchvision.transforms.ToTensor())
validation_dataset = torchvision.datasets.CIFAR10("./cifar10_data", train=False, download=True,
                                            transform=torchvision.transforms.ToTensor())

データセットのサイズ

training_dataset_size = len(training_dataset)
validation_dataset_size = len(validation_dataset)
print("トレーニングセットのサイズ：{}".format(training_dataset_size))
print("検証セットのサイズ：{}".format(validation_dataset_size))

DataLoaderを使用したデータセットの読み込み

from torch.utils.data import DataLoader

training_loader = DataLoader(training_dataset, batch_size=64, shuffle=True)
validation_loader = DataLoader(validation_dataset, batch_size=64, shuffle=False)

ニューラルネットワークの構築

import torch.nn as nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear

class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.feature_extractor = Sequential(
            Conv2d(3, 32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            MaxPool2d(kernel_size=2),
            Conv2d(32, 32, kernel_size=5, padding=2),
            nn.ReLU(),
            MaxPool2d(kernel_size=2),
            Conv2d(32, 64, kernel_size=5, padding=2),
            nn.ReLU(),
            MaxPool2d(kernel_size=2),
            Flatten(),
            Linear(1024, 64),
            nn.ReLU(),
            Linear(64, 10)
        )

    def forward(self, x):
        x = self.feature_extractor(x)
        return x

ネットワークモデルの作成

my_model = ImageClassifier()

損失関数の定義

criterion = nn.CrossEntropyLoss()

最適化アルゴリズムの設定

learning_rate = 0.01
optimizer = torch.optim.SGD(my_model.parameters(), lr=learning_rate)

トレーニングのパラメータ設定

total_train_steps = 0
total_validation_steps = 0
num_epochs = 10

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("./tensorboard_logs")

トレーニングと検証

for epoch_num in range(num_epochs):
    print(f"----------エポック {epoch_num + 1} のトレーニングを開始----------")

    # トレーニングステップ
    my_model.train()
    for batch in training_loader:
        images, labels = batch
        predictions = my_model(images)
        current_loss = criterion(predictions, labels)

        optimizer.zero_grad()
        current_loss.backward()
        optimizer.step()

        total_train_steps += 1
        if total_train_steps % 100 == 0:
            print(f"トレーニングステップ: {total_train_steps}, 損失: {current_loss.item():.4f}")
            writer.add_scalar("Loss/Train", current_loss.item(), total_train_steps)

    # 検証ステップ
    my_model.eval()
    total_validation_loss = 0.0
    correct_predictions = 0
    with torch.no_grad():
        for batch in validation_loader:
            images, labels = batch
            predictions = my_model(images)
            loss = criterion(predictions, labels)
            total_validation_loss += loss.item()
            correct_predictions += (predictions.argmax(1) == labels).type(torch.float).sum().item()

    avg_validation_loss = total_validation_loss / len(validation_loader)
    validation_accuracy = correct_predictions / validation_dataset_size
    print(f"検証セット全体の平均損失: {avg_validation_loss:.4f}")
    print(f"検証セット全体の正解率: {validation_accuracy:.4f}")
    writer.add_scalar("Loss/Validation", avg_validation_loss, epoch_num)
    writer.add_scalar("Accuracy/Validation", validation_accuracy, epoch_num)
    total_validation_steps += 1

    # モデルの保存
    torch.save(my_model.state_dict(), f"classifier_epoch_{epoch_num}.pth")
    print("モデルを保存しました")

writer.close()