データセット準備

import torch
import torch.nn as nn
from torchvision import transforms, datasets

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

# データ変換処理
training_transforms = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# データセット読み込み
weather_data = datasets.ImageFolder("./data/weather_photos/", transform=training_transforms)
class_labels = weather_data.classes

# データ分割
train_size = int(0.8 * len(weather_data))
test_size = len(weather_data) - train_size
train_set, test_set = torch.utils.data.random_split(weather_data, [train_size, test_size])

# データローダー設定
batch_size = 8
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=True)

ネットワークコンポーネント実装

class StandardConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel=1, stride=1, padding=None, groups=1, activation=True):
        super().__init__()
        self.conv_layer = nn.Conv2d(in_channels, out_channels, kernel, stride, self.auto_pad(kernel, padding), groups=groups, bias=False)
        self.batch_norm = nn.BatchNorm2d(out_channels)
        self.activation = nn.SiLU() if activation else nn.Identity()
    
    def auto_pad(self, kernel, padding):
        return kernel // 2 if padding is None else padding
    
    def forward(self, x):
        return self.activation(self.batch_norm(self.conv_layer(x)))

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, shortcut=True, groups=1, expansion=0.5):
        super().__init__()
        hidden_channels = int(out_channels * expansion)
        self.conv1 = StandardConv(in_channels, hidden_channels, 1)
        self.conv2 = StandardConv(hidden_channels, out_channels, 3, groups=groups)
        self.use_shortcut = shortcut and in_channels == out_channels
    
    def forward(self, x):
        return x + self.conv2(self.conv1(x)) if self.use_shortcut else self.conv2(self.conv1(x))

class C3Block(nn.Module):
    def __init__(self, in_channels, out_channels, iterations=1, shortcut=True, groups=1, expansion=0.5):
        super().__init__()
        hidden_channels = int(out_channels * expansion)
        self.conv1 = StandardConv(in_channels, hidden_channels, 1)
        self.conv2 = StandardConv(in_channels, hidden_channels, 1)
        self.conv3 = StandardConv(2 * hidden_channels, out_channels, 1)
        self.residual_blocks = nn.Sequential(*(ResidualBlock(hidden_channels, hidden_channels, shortcut, groups) for _ in range(iterations)))
    
    def forward(self, x):
        return self.conv3(torch.cat((self.residual_blocks(self.conv1(x)), self.conv2(x)), dim=1))

class FastSPP(nn.Module):
    def __init__(self, in_channels, out_channels, kernel=5):
        super().__init__()
        mid_channels = in_channels // 2
        self.conv1 = StandardConv(in_channels, mid_channels, 1)
        self.conv2 = StandardConv(mid_channels * 4, out_channels, 1)
        self.pool = nn.MaxPool2d(kernel_size=kernel, stride=1, padding=kernel // 2)
    
    def forward(self, x):
        x = self.conv1(x)
        y1 = self.pool(x)
        y2 = self.pool(y1)
        y3 = self.pool(y2)
        return self.conv2(torch.cat([x, y1, y2, y3], 1))

バックボーンネットワーク

class YOLOv5Backbone(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            StandardConv(3, 64, 3, 2, 2),
            StandardConv(64, 128, 3, 2),
            C3Block(128, 128),
            StandardConv(128, 256, 3, 2),
            C3Block(256, 256),
            StandardConv(256, 512, 3, 2),
            C3Block(512, 512),
            StandardConv(512, 1024, 3, 2),
            C3Block(1024, 1024),
            FastSPP(1024, 1024, 5)
        )
        self.classifier = nn.Sequential(
            nn.Linear(1024 * 8 * 8, 100),
            nn.ReLU(),
            nn.Linear(100, len(class_labels))
        )
    
    def forward(self, x):
        x = self.layers(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)

model = YOLOv5Backbone().to(device)

トレーニングプロセス

def train_model(loader, model, loss_fn, optimizer):
    model.train()
    total_loss, total_correct = 0, 0
    for images, targets in loader:
        images, targets = images.to(device), targets.to(device)
        outputs = model(images)
        loss = loss_fn(outputs, targets)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        total_correct += (outputs.argmax(1) == targets).sum().item()
    
    return total_correct / len(loader.dataset), total_loss / len(loader)

def evaluate_model(loader, model, loss_fn):
    model.eval()
    total_loss, total_correct = 0, 0
    with torch.no_grad():
        for images, targets in loader:
            images, targets = images.to(device), targets.to(device)
            outputs = model(images)
            total_loss += loss_fn(outputs, targets).item()
            total_correct += (outputs.argmax(1) == targets).sum().item()
    
    return total_correct / len(loader.dataset), total_loss / len(loader)

# トレーニング設定
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_function = nn.CrossEntropyLoss()
epochs = 30

for epoch in range(epochs):
    train_acc, train_loss = train_model(train_loader, model, loss_function, optimizer)
    test_acc, test_loss = evaluate_model(test_loader, model, loss_function)
    print(f'Epoch:{epoch+1:2d}, Train Acc:{train_acc*100:.1f}%, Train Loss:{train_loss:.3f}, '
          f'Test Acc:{test_acc*100:.1f}%, Test Loss:{test_loss:.3f}')

ネットワークモジュール解説

Focusモジュール

高解像度画像からピクセルを抽出し、特徴マップの次元を変換。入力画像をスライスして連結することで、空間情報をチャネル次元に集約。

標準畳み込み

2D畳み込み層＋バッチ正規化＋SiLU活性化関数の組み合わせ。計算効率と精度のバランスを最適化。

C3ブロック

CSP（Cross Stage Partial）アーキテクチャを採用。複数のResidualBlockを並列処理し、特徴抽出効率を向上。

高速空間ピラミッドプーリング

複数スケールの最大プーリング操作を統合。異なる受容野の特徴を効率的に抽出し、オブジェクトスケールの変化に頑健。