Pytorch Lightning – Auto Encoder で MNIST の特徴表現を学習する

from pathlib import Path
from typing import Optional, Union, Tuple

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 MNISTDataModule(pl.LightningDataModule):
    def __init__(self, batch_size, data_dir="./"):
        super().__init__()

        dataset = torchvision.datasets.MNIST(
            data_dir, train=True, download=True, transform=transforms.ToTensor()
        )
        self.train_dataset, self.val_dataset = torch.utils.data.random_split(
            dataset, [55000, 5000]
        )
        self.test_dataset = torchvision.datasets.MNIST(
            data_dir, train=False, download=True, transform=transforms.ToTensor()
        )

        self.batch_size = batch_size
        self.data_dir = data_dir

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

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

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

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

class LitAutoEncoder(pl.LightningModule):
    def __init__(self, hidden_dim=64):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 32)
        )
        self.decoder = nn.Sequential(
            nn.Linear(32, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 28 * 28)
        )

    def forward(self, x):
        z = self.encoder(x)
        x_hat = self.decoder(z)
        return x_hat

    def training_step(self, batch, batch_idx):
        return self._common_step(batch, batch_idx, "train")

    def validation_step(self, batch, batch_idx):
        self._common_step(batch, batch_idx, "val")

    def test_step(self, batch, batch_idx):
        self._common_step(batch, batch_idx, "test")

    def predict_step(self, batch, batch_idx, dataloader_idx=None):
        x, _ = batch
        return self(x)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

    def _prepare_batch(self, batch):
        x, _ = batch
        return x.view(x.size(0), -1)

    def _common_step(self, batch, batch_idx, stage):
        x = self._prepare_batch(batch)
        loss = F.mse_loss(x, self(x))
        self.log(f"{stage}_loss", loss, on_step=True)
        return loss

class ImageSampler(pl.callbacks.Callback):
    def __init__(
        self, output_dir: str = "./generated", num_samples: int = 64, nrow: int = 8
    ):
        super().__init__()
        output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True)

        self.output_dir = output_dir
        self.num_samples = num_samples
        self.nrow = nrow

    def on_epoch_end(self, trainer: pl.Trainer, model: pl.LightningModule) -> None:
        # 生成する枚数をバッチサイズに指定し、検証用のデータセットから取得する。
        dataloader = torch.utils.data.DataLoader(
            trainer.datamodule.val_dataset, batch_size=self.num_samples
        )
        img_in, _ = next(iter(dataloader))

        x = img_in.view(img_in.size(0), -1)  # (N, C, H, W) -> (N, C * H * W)

        with torch.no_grad():
            model.eval()
            y = model(x.to(model.device))
            model.train()

        img_out = y.reshape(img_in.shape)  # (N, C * H * W) -> (N, C, H, W)

        # 画像を結合する。
        img_grid = torchvision.utils.make_grid(img_out, nrow=self.nrow, normalize=True)

        # 画像を保存する。
        save_path = self.output_dir / f"generated_epoch{trainer.current_epoch}.png"
        torchvision.utils.save_image(img_grid, save_path)

dm = MNISTDataModule(batch_size=32, data_dir="/data/MNIST/")
model = LitAutoEncoder()

trainer = pl.Trainer(
    max_epochs=10, log_every_n_steps=1, gpus=1, callbacks=[ImageSampler()]
)
trainer.fit(model, dm)

def extract(dataloader):
    X, y = [], []
    for x_batch, y_batch in dataloader:
        x_batch = x_batch.view(x_batch.size(0), -1)
        feat = model.encoder(x_batch)

        X.append(feat)
        y.append(y_batch)

    X = torch.concat(X).detach().numpy()
    y = torch.concat(y).detach().numpy()

    return X, y


# 学習データ、テストデータを Encoder で特徴量に変換する
X_train, y_train = extract(dm.train_dataloader())
X_test, y_test = extract(dm.test_dataloader())

from sklearn import metrics, svm

clf = svm.SVC(kernel="linear", C=1, random_state=0)
clf.fit(X_train, y_train)

y_pred = clf.predict(X_test)

acc = metrics.accuracy_score(y_test, y_pred)
print(f"accuracy: {acc:.2%}")