EfficientNet-B0で麻雀牌を学習する｜PyTorch・timmで転移学習

import timm
import torch
from torch import nn
from tqdm import tqdm
from pathlib import Path
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from sklearn.metrics import accuracy_score



class ImageClassifierTrainer:
    def __init__(self, train_dataset, val_dataset, test_dataset, output_model,
                 image_size=224, batch_size=32,
                 epochs=3, lr=1e-4, device="mps",
                 model_name="efficientnet_b0"):
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.test_dataset = test_dataset

        self.class_names = train_dataset.class_names

        self.num_classes = len(self.class_names)
        
        pin_memory = device.startswith("cuda")

        self.train_loader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=True,
            pin_memory=pin_memory,
        )

        self.val_loader = DataLoader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            pin_memory=pin_memory,
        )

        self.test_loader = DataLoader(
            test_dataset,
            batch_size=batch_size,
            shuffle=False,
            pin_memory=pin_memory,
        )
        self.output_model=Path(output_model)
        self.image_size=image_size
        self.batch_size=batch_size
        self.epochs=epochs
        self.lr=lr
        self.device=device
        self.train_tf=transforms.Compose([
            transforms.Resize((image_size,image_size)),
            transforms.RandomRotation(2),
            transforms.ColorJitter(0.15,0.15,0.15),
            transforms.ToTensor(),
            transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
        ])
        self.val_tf=transforms.Compose([
            transforms.Resize((image_size,image_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
        ])
        
        self.train_dataset.transform = self.train_tf
        self.val_dataset.transform = self.val_tf
        self.test_dataset.transform = self.val_tf

        self.model = timm.create_model(model_name, pretrained=True, num_classes=self.num_classes,).to(device)
        
        self.criterion=nn.CrossEntropyLoss()
        self.optimizer=torch.optim.AdamW(self.model.parameters(),lr=self.lr)
        self.scheduler=torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer,T_max=self.epochs)

    def _run_epoch(self, train=True):
        self.model.train(train)
        loader=self.train_loader if train else self.val_loader
        preds=[]; labels=[]; loss_sum=0
        ctx=torch.enable_grad() if train else torch.no_grad()
        with ctx:
            for x,y in tqdm(loader, leave=False):
                x,y=x.to(self.device),y.to(self.device)
                if train: self.optimizer.zero_grad()
                out=self.model(x)
                loss=self.criterion(out,y)
                if train:
                    loss.backward(); self.optimizer.step()
                loss_sum+=loss.item()
                preds.extend(out.argmax(1).cpu().tolist())
                labels.extend(y.cpu().tolist())
        return loss_sum/len(loader), accuracy_score(labels,preds)

    def train(self):
        best=0
        for e in range(self.epochs):
            tl,ta=self._run_epoch(True)
            vl,va=self._run_epoch(False)
            self.scheduler.step()
            print(f"Epoch {e+1}: train {ta:.4f} val {va:.4f}")
            if va>best:
                best=va
                self.output_model.parent.mkdir(parents=True,exist_ok=True)
                torch.save({"model_state_dict":self.model.state_dict(),
                            "class_names":self.class_names}, self.output_model)

        print("Best:",best)
        if self.test_dataset is not None:
            self.test()

    @torch.no_grad()
    def test(self):
        self.model.eval()
        rows = []
        correct = 0
        for batch_idx, (images, labels) in enumerate(self.test_loader):
            images = images.to(self.device)
            outputs = self.model(images)
            probs = F.softmax(outputs, dim=1)
            preds = probs.argmax(dim=1).cpu()
            labels = labels.cpu()

            for i in range(len(preds)):
                pred_idx = preds[i].item()
                label_idx = labels[i].item()
                gt = self.class_names[label_idx]
                pred = self.class_names[pred_idx]
                ok = gt == pred
                if ok:
                    correct += 1

                sample_idx = batch_idx * self.batch_size + i
                image_path = str(
                    self.test_dataset.samples[sample_idx][0]
                )
                row = {
                    "image": image_path,
                    "ground_truth": gt,
                    "prediction": pred,
                    "correct": ok,
                }

                for c, class_name in enumerate(self.class_names):
                    row[f"prob_{class_name}"] = float(probs[i, c])

                rows.append(row)

        accuracy = correct / len(self.test_dataset)
        print(f"Test Accuracy : {accuracy:.4f}")