Pytorch – Vanilla Backpropagation で顕著正マップを作成する方法について

import json
from collections import defaultdict
from inspect import isfunction
from pathlib import Path

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision
from opencv_transforms import transforms
from torch.nn import functional as F
from torch.utils.data import DataLoader, Dataset
from torchvision.datasets.utils import download_url

def get_device(use_gpu):
    if use_gpu and torch.cuda.is_available():
        # これを有効にしないと、計算した勾配が毎回異なり、再現性が担保できない。
        torch.backends.cudnn.deterministic = True
        return torch.device("cuda")
    else:
        return torch.device("cpu")


# デバイスを選択する。
device = get_device(use_gpu=True)

def _get_avaable_models():
    available_models = {}
    for obj_name in dir(torchvision.models):
        obj = getattr(torchvision.models, obj_name)
        if isfunction(obj):
            available_models[obj_name] = obj

    return available_models

def get_model(name, device):
    available_models = _get_avaable_models()
    if name not in available_models:
        raise ValueError(f"{name} という名前のモデルは torchvision に存在しません。")

    model_fn = available_models[name]
    model = model_fn(pretrained=True).to(device)
    model.eval()

    return model


# モデルを作成する。
model = get_model("alexnet", device)

def _get_img_paths(img_dir):
    img_dir = Path(img_dir)

    img_extensions = [".jpg", ".jpeg", ".png", ".bmp"]
    img_paths = [str(p) for p in img_dir.iterdir() if p.suffix in img_extensions]
    img_paths.sort()

    return img_paths

class ImageFolder(Dataset):
    IMAGENET_SIZE = (224, 224)
    IMAGENET_MEAN = [0.485, 0.456, 0.406]
    IMAGENET_STD = [0.229, 0.224, 0.225]

    def __init__(self, img_dir):
        # 画像ファイルのパス一覧を取得する。
        self.img_paths = _get_img_paths(img_dir)

        # Transformer を作成する。
        self.transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=ImageFolder.IMAGENET_MEAN, std=ImageFolder.IMAGENET_STD
                ),
            ]
        )

    def __getitem__(self, index):
        path = self.img_paths[index]

        # 画像を読み込む。
        raw_img = cv2.imread(path)

        # チャンネル順を BGR から RGB に変換する。
        raw_img = cv2.cvtColor(raw_img, cv2.COLOR_BGR2RGB)

        # リサイズする。
        raw_img = cv2.resize(raw_img, ImageFolder.IMAGENET_SIZE)

        # Transformer で変換する。
        img = self.transform(raw_img)

        sample = {"raw_image": raw_img, "image": img, "path": path}

        return sample

    def __len__(self):
        return len(self.img_paths)


# Dataset を作成する。
dataset = ImageFolder("data")

# DataLoader を作成する。
dataloader = DataLoader(dataset, batch_size=4)

def get_classes():
    if not Path("data/imagenet_class_index.json").exists():
        # ファイルが存在しない場合はダウンロードする。
        download_url("https://git.io/JebAs", "data", "imagenet_class_index.json")

    # クラス一覧を読み込む。
    with open("data/imagenet_class_index.json") as f:
        data = json.load(f)
        class_names = [x["ja"] for x in data]

    return class_names


# クラス名一覧を取得する。
class_names = get_classes()

def to_onehot(y, n_classes):
    # (N,) tensor to (N, NumClasses) tensor.
    return torch.eye(n_classes)[y]


class VanillaBackprop:
    def __init__(self, model):
        self.model = model
        self.device = next(model.parameters()).device

    def forward(self, inputs):
        inputs = inputs.to(self.device).requires_grad_()

        # 勾配を初期化する。
        self.model.zero_grad()

        # 順伝搬を行う。
        logits = self.model(inputs)

        # softmax を適用する。
        probs = F.softmax(logits, dim=1)

        # 確率が大きい順にソートし、対応するクラス ID の一覧を返す。
        class_ids = probs.sort(dim=1, descending=True)[1]

        self.inputs = inputs
        self.logits = logits

        return class_ids

    def backward(self, class_ids):
        # 逆伝搬の入力を作成する。
        onehot = to_onehot(class_ids, n_classes=self.logits.size(1)).to(self.device)

        # 逆伝搬を行う。
        self.logits.backward(gradient=onehot, retain_graph=True)

    def generate(self):
        # 入力層の勾配を取得する。
        gradients = self.inputs.grad.clone()
        # 入力層の勾配を初期化する。
        self.inputs.grad.zero_()

        return gradients

def normalize(x):
    B, C, H, W = x.shape
    x = x.view(B, -1)
    xmin = x.min(dim=1, keepdim=True)[0]
    xmax = x.max(dim=1, keepdim=True)[0]
    x = (x - xmin) / (xmax - xmin)
    x = x.view(B, C, H, W)

    return x


def tensor_to_images(x):
    x = normalize(x)
    x = x.cpu().numpy().transpose(0, 2, 3, 1)
    x = (x * 255).astype(np.uint8)

    if x.shape[3] == 1:
        # (B, H, W, 1) -> (B, H, W)
        x = x.squeeze(axis=3)

    return x

def tensor_to_gray_images(x):
    x = x.abs().sum(dim=1, keepdims=True)
    x = tensor_to_images(x)

    return x

def tensor_to_positive_saliency(x):
    x = torch.nn.functional.relu(x)
    x = tensor_to_images(x)

    return x

def tensor_to_negative_saliency(x):
    x = torch.nn.functional.relu(-x)
    x = tensor_to_images(x)

    return x

# Vanilla Backpropagation
vanilla_backprop = VanillaBackprop(model)

results = defaultdict(list)
for batch in dataloader:
    # 順伝搬を行う。
    class_ids = vanilla_backprop.forward(batch["image"])

    for i in range(3):  # top 1 ~ top 3
        # i 番目にスコアが高いクラスを取得する。
        ith_class_ids = class_ids[:, i]

        # 逆伝搬を行う。
        vanilla_backprop.backward(ith_class_ids)

        # 入力画像に対する勾配を取得する。
        gradients = vanilla_backprop.generate()

        # 勾配を画像にする。
        color_grad_imgs = tensor_to_images(gradients)
        gray_grad_imgs = tensor_to_gray_images(gradients)
        pos_grad_imgs = tensor_to_positive_saliency(gradients)
        neg_grad_imgs = tensor_to_negative_saliency(gradients)

        for j in range(len(batch["image"])):
            img = batch["raw_image"][j].numpy()
            img_path = batch["path"][j]
            class_id = ith_class_ids[j]
            color_grad = color_grad_imgs[j]
            gray_grad = gray_grad_imgs[j]
            pos_grad = pos_grad_imgs[j]
            neg_grad = neg_grad_imgs[j]

            result = {
                "image": img,
                "color_grad": color_grad,
                "gray_grad": gray_grad,
                "pos_grad": pos_grad,
                "neg_grad": neg_grad,
                "class_id": class_id,
                "class_name": class_names[class_id],
            }

            results[img_path].append(result)

def output_results(output_dir, results, name):
    output_dir.mkdir(exist_ok=True)

    for img_path, result in results.items():
        fig = plt.figure(figsize=(10, 10))

        ax = fig.add_subplot(4, 4, 1)
        ax.set_axis_off()
        ax.imshow(result[0]["image"])

        # top1 ~ top3 スコアのクラスの勾配を描画する。
        for i in range(3):
            color_img = result[i]["color_grad"]
            gray_img = result[i]["gray_grad"]
            pos_grad = result[i]["pos_grad"]
            neg_grad = result[i]["neg_grad"]
            class_name = result[i]["class_name"]

            ax = fig.add_subplot(4, 4, i * 4 + 5)
            ax.imshow(color_img)
            ax.set_ylabel(f"Top {i + 1} - {class_name}")
            ax.set_xticks([])
            ax.set_yticks([])

            ax = fig.add_subplot(4, 4, i * 4 + 6)
            ax.imshow(gray_img, cmap="jet")
            ax.set_axis_off()

            ax = fig.add_subplot(4, 4, i * 4 + 7)
            ax.imshow(pos_grad, cmap="jet")
            ax.set_axis_off()

            ax = fig.add_subplot(4, 4, i * 4 + 8)
            ax.imshow(neg_grad, cmap="jet")
            ax.set_axis_off()

        axes = fig.get_axes()
        axes[1].set_title("Color Gradient")
        axes[2].set_title("Grayscale Gradient")
        axes[3].set_title("Positive Gradient")
        axes[4].set_title("Negative Gradient")

        # ax.set_title(f"top {i + 1} {class_name}", fontsize=10)
        save_path = output_dir / f"{name}_{Path(img_path).stem}.png"
        fig.savefig(save_path, bbox_inches="tight")


# 結果を保存する。
output_results(Path("output"), results, "vanilla_backprop")