画像分類モデルをゼロから構築してみた
– AIデータセンター活用記

[ログインサーバー]$ vi ~/.config/enroot/.credentials

# 以下を追記して保存

machine nvcr.io login $oauthtoken password <Personal API Key>
machine authn.nvidia.com login $oauthtoken password <Personal API Key>

[ログインサーバ]$ cd /sample
[ログインサーバ]$ enroot import docker://nvcr.io/nvidia/pytorch:24.03-py3

[ログインサーバー]$
srun -N1 -n1 --gpus=1 --time=00:30:00 \
-p <partition> \
--container-image=/sample/nvidia+pytorch+24.03-py3.sqsh \
--container-mounts=/sample:/sample \
--pty bash

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

# CIFAR-10データをダウンロード
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)

# シンプルなCNN
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64*8*8, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.fc(self.conv(x))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 学習ループ
for epoch in range(3):  # 3エポックだけ
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 100 == 99:
            print(f"[{epoch+1}, {i+1:5d}] loss: {running_loss / 100:.3f}")
            running_loss = 0.0
print("Finished Training")
torch.save(net.state_dict(), "/lustre/sample/cifar_net.pth")
print(" モデルを保存しました。")

[ログインサーバ]$ python cifar_train.py

# 出力例
[1, 100] loss: 1.957
[1, 200] loss: 1.606
・
・（省略）
・
[3, 1500] loss: 0.716

Finished Training
モデルを保存しました。

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn

# CIFAR-10のテストデータ
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=32, shuffle=False)

# モデル定義（train.pyと同じにする必要あり）
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64*8*8, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.fc(self.conv(x))

# デバイス設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = SimpleCNN().to(device)

# ② モデル読み込み
net.load_state_dict(torch.load("/lustre/sample/cifar_net.pth"))
net.eval()

# 推論と評価
correct = 0
total = 0
with torch.no_grad():
    for i, data in enumerate(testloader):  # ← enumerateを追加
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

        # ここに進捗表示を追加

        if (i + 1) % 100 == 0:
　　　　print(f"[{i+1}] 現在の精度: {100 * correct / total:.2f}%", flush=True)

print(f" テスト精度: {100 * correct / total:.2f}%")

userXXX@server:/sample$ python cifar_test.py

# 出力例

[100]  現在の精度： 71.56%
[200]  現在の精度： 71.28%
[300]  現在の精度： 71.72%
テスト精度: 71.65%