たった 3 行で性能向上！torchvision のデータ拡張で AI を学習させてみた！

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from torchvision.models import vgg16

!pip install pytorch-lightning

import torchvision
from torchvision import transforms, datasets
import pytorch_lightning as pl

# データセットの変換を定義
DA_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.CenterCrop(24),
    transforms.Resize(32), 
    transforms.ToTensor()
])

# データセットの取得
train_val = datasets.CIFAR10('./', train=True, download=True, transform=DA_transform)
test = datasets.CIFAR10('./', train=False, download=True, transform=transforms.Totensor())

# train と val に分割
pl.seed_everything(0)
n_train, n_val = int(len(train_val)*0.8), int(len(train_val)*0.2)
train, val = torch.utils.data.random_split(train_val, [n_train, n_val])

# バッチサイズの定義
batch_size = 256

# Data Loader を定義
train_loader = torch.utils.data.DataLoader(train, batch_size, shuffle=True, drop_last=True, num_workers=4)
val_loader = torch.utils.data.DataLoader(val, batch_size, num_workers=4)
test_loader = torch.utils.data.DataLoader(test, batch_size)

# 学習データの先頭 100 サンプルを確認
plt.figure(figsize=(10, 10))
for n in range(16):
    x, t = train[n]
    plt.subplot(4, 4, n+1)
    plt.title(train_val.classes[train[n][1]])
    plt.axis('off')
    plt.imshow(x.permute(1, 2, 0))

# 評価指標の設定
from torchmetrics.classification import MulticlassAccuracy
accuracy = MulticlassAccuracy(num_classes=10).to(device='cuda')

# モデルの設計
class CnnClassifier(pl.LightningModule):

    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.fc = nn.Linear(64*8*8, 10)

    def forward(self, x):
        h = self.conv1(x)
        h = F.relu(h)
        h = self.bn1(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = self.conv2(h)
        h = F.relu(h)
        h = self.bn2(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = h.view(-1, 64*8*8)
        h = self.fc(h)
        return h

    def training_step(self, batch, batch_idx):
        x, t = batch
        y = self(x)
        loss = F.cross_entropy(y, t)
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        self.log('train_acc', accuracy(y, t), on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, t = batch
        y = self(x)
        loss = F.cross_entropy(y, t)
        self.log('val_loss', loss, on_step=False, on_epoch=True)
        self.log('val_acc', accuracy(y, t), on_step=False, on_epoch=True)
        return loss

    def test_step(self, batch, batch_idx):
        x, t = batch
        y = self(x)
        loss = F.cross_entropy(y, t)
        self.log('test_loss', loss, on_step=False, on_epoch=True)
        self.log('test_acc', accuracy(y, t), on_step=False, on_epoch=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.01)
        return optimizer

# 学習の実行
net = CnnClassifier() # 定義したモデルの呼び出し
trainer = pl.Trainer(max_epochs=20, accelerator='gpu') # 学習方法の設定
trainer.fit(net, train_loader, val_loader) # モデルを学習

# テストデータで検証
results = trainer.test(dataloaders=test_loader)
results