【AI Shift Advent Calendar 2021】DNNのモデル分析ツールWeight Watcherを使ってみる

import os 
import random
import wandb
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import weightwatcher as ww
from torchvision import datasets, transforms

class GCF:
    EXP_NAME = 'exp_01'
    LR = 1e-3
    BS = 32
    N_EPOCH = 20
    SEED = 0
    MLP_UNIT = 64

class MnistClassifier(nn.Module):
    def __init__(self):
        super(MnistClassifier, self).__init__()
        self.layer1 = nn.Linear(28 * 28, GCF.MLP_UNIT)
        self.activation1 = nn.ReLU()
        self.layer2 = nn.Linear(GCF.MLP_UNIT, GCF.MLP_UNIT)
        self.activation2 = nn.ReLU()
        self.layer3 = nn.Linear(GCF.MLP_UNIT, 10)
    def forward(self, X):
        h = self.layer1(X.view(-1, 28 * 28))
        h = self.activation1(h)
        h = self.layer2(h)
        h = self.activation2(h)
        h = self.layer3(h)
        return h

mnist_data = datasets.MNIST('data', train=True, download=True, transform=transforms.ToTensor())
mnist_data = list(mnist_data)
mnist_train = mnist_data[:1000]
mnist_valid = mnist_data[1000:2000]

train_loader = torch.utils.data.DataLoader(mnist_train,
                                          batch_size=GCF.BS,
                                          shuffle=True)
valid_loader = torch.utils.data.DataLoader(mnist_valid,
                                          batch_size=GCF.BS,
                                          shuffle=False)

def train_loop(model, train_loader):
    model.train()
    for imgs, labels in train_loader:
        out = model(imgs)
        loss = criterion(out, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

def valid_loop(model, valid_loader):
    model.eval()
    valid_losses, valid_predicts = [], []
    for imgs, labels in valid_loader:
        with torch.no_grad():
            out = model(imgs)
        loss = criterion(out, labels)
        valid_losses.append(loss.item())
        valid_predicts.append(out.argmax(1))
        
    valid_loss = np.array(valid_losses).mean()
    valid_predict = torch.hstack(valid_predicts).numpy()
    y_valid = np.array([j for i, j in valid_loader.dataset])

    valid_acc = (y_valid == valid_predict).mean()

    return {
        'valid_loss': valid_loss,
        'valid_acc': valid_acc,
    }

model = MnistClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=GCF.LR)

# 学習するモデルをWeightWatcherにわたす
watcher = ww.WeightWatcher(model=model)

for epoch in range(GCF.N_EPOCH):
    train_loop(model, train_loader)
    res = valid_loop(model, valid_loader)

    # モデルの分析
    details = watcher.analyze(randomize=True)
    # Rand Distanceの計算(各層の値が出るので平均をとる)
    avg_rand_distance = details.rand_distance.mean()

    res.update({'rand_distance': avg_rand_distance})
    print(res)