Pytorch Lightning – 事前学習モデルを使ってクラス分類モデルを学習する方法

from pathlib import Path
from typing import Optional

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

class HymenopteraDataModule(pl.LightningDataModule):
    def __init__(self, batch_size: int, data_dir: str = "./"):
        super().__init__()
        self.data_dir = Path(data_dir)
        self.batch_size = batch_size
        self.train_transform = self._get_transform(train=True)
        self.test_transform = self._get_transform(train=False)

    def prepare_data(self):
        # ダウンロードして、解凍する。
        url = "https://download.pytorch.org/tutorial/hymenoptera_data.zip"
        torchvision.datasets.utils.download_and_extract_archive(url, self.data_dir)

    def setup(self, stage: Optional[str] = None):
        if stage == "fit" or stage is None:
            # train/ の画像は学習、検証用に使用する。
            dataset = torchvision.datasets.ImageFolder(
                self.data_dir / "hymenoptera_data/train"
            )

            # train/ の画像を学習用と検証用に分割する。
            n_train = int(len(dataset) * 0.8)
            n_val = len(dataset) - n_train
            self.train_dataset, self.val_dataset = torch.utils.data.random_split(
                dataset, [n_train, n_val]
            )
            self.train_dataset.dataset.transform = self.train_transform
            self.val_dataset.dataset.transform = self.test_transform
            self.classes = self.train_dataset.dataset.classes

        if stage == "test" or stage is None:
            # val/ の画像はテスト用に使用する。
            self.test_dataset = torchvision.datasets.ImageFolder(
                self.data_dir / "hymenoptera_data/val"
            )
            self.test_dataset.transform = self.test_transform
            self.classes = self.test_dataset.classes

    def train_dataloader(self):
        return torch.utils.data.DataLoader(
            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4
        )

    def val_dataloader(self):
        return torch.utils.data.DataLoader(
            self.val_dataset, batch_size=self.batch_size, num_workers=4
        )

    def test_dataloader(self):
        return torch.utils.data.DataLoader(
            self.test_dataset, batch_size=self.batch_size, num_workers=4
        )

    def predict_dataloader(self):
        return torch.utils.data.DataLoader(
            self.test_dataset, batch_size=self.batch_size, num_workers=4
        )

    def _get_transform(self, train: bool):
        lst = []
        if train:
            # 学習時に適用する Transform
            lst += [
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
            ]
        else:
            # 推論時に適用する Transform
            lst += [
                transforms.Resize(256),
                transforms.CenterCrop(224),
            ]

        lst += [
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]

        transform = transforms.Compose(lst)

        return transform

class LitModel(pl.LightningModule):
    def __init__(self, num_classes):
        super().__init__()
        # ResNet-18 を作成する。
        self.model = torchvision.models.resnet18(pretrained=True)

        # 出力層の出力数を ImageNet の 1000 からこのデータセットのクラス数である 2 に置き換える。
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        loss, acc = self._shared_step(batch)

        self.log("train_loss", loss)
        self.log("train_acc", acc)

        return loss

    def validation_step(self, batch, batch_idx):
        loss, acc = self._shared_step(batch)

        self.log("val_loss", loss)
        self.log("val_acc", acc)

        return loss

    def test_step(self, batch, batch_idx):
        loss, acc = self._shared_step(batch)

        self.log("test_loss", loss, prog_bar=True)
        self.log("test_acc", acc, prog_bar=True)

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        x, _ = batch
        y_hat = self.model(x)
        preds = torch.argmax(y_hat, dim=1)

        return preds

    def configure_optimizers(self):
        optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

        return [optimizer], [scheduler]

    def _shared_step(self, batch):
        x, y = batch
        y_hat = self.model(x)
        loss = F.cross_entropy(y_hat, y)

        # accuracy を計算する。
        preds = torch.argmax(y_hat, dim=1)
        acc = (y == preds).sum() / y.size(0)

        return loss, acc

dm = HymenopteraDataModule(batch_size=8, data_dir="/data/hymenoptera_data/")
model = LitModel(num_classes=2)

trainer = pl.Trainer(max_epochs=10, log_every_n_steps=1, gpus=1)

# 学習する。
# !rm -r lightning_logs
trainer.fit(model, dm)

%reload_ext tensorboard
%tensorboard --logdir ./lightning_logs --bind_all --port 6006

ret = trainer.test(datamodule=dm, ckpt_path="best")