preparation.py

# PyTorch StudioGAN: https://github.com/POSTECH-CVLab/PyTorch-StudioGAN
# The MIT License (MIT)
# See license file or visit https://github.com/POSTECH-CVLab/PyTorch-StudioGAN for details

# src/metrics/preparation.py

from os.path import exists, join
import os

try:
    from torchvision.models.utils import load_state_dict_from_url
except ImportError:
    from torch.utils.model_zoo import load_url as load_state_dict_from_url
from torch.nn import DataParallel
from torch.nn.parallel import DistributedDataParallel as DDP
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np

from metrics.inception_net import InceptionV3
from metrics.swin_transformer import SwinTransformer
import metrics.features as features
import metrics.vit as vits
import metrics.fid as fid
import metrics.ins as ins
import utils.misc as misc
import utils.ops as ops
import utils.resize as resize

model_versions = {"InceptionV3_torch": "pytorch/vision:v0.10.0",
                  "ResNet_torch": "pytorch/vision:v0.10.0",
                  "SwAV_torch": "facebookresearch/swav:main"}
model_names = {"InceptionV3_torch": "inception_v3",
               "ResNet50_torch": "resnet50",
               "SwAV_torch": "resnet50"}
SWAV_CLASSIFIER_URL = "https://dl.fbaipublicfiles.com/deepcluster/swav_800ep_eval_linear.pth.tar"
SWIN_URL = "https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22kto1k.pth"


class LoadEvalModel(object):
    def __init__(self, eval_backbone, post_resizer, world_size, distributed_data_parallel, device):
        super(LoadEvalModel, self).__init__()
        self.eval_backbone = eval_backbone
        self.post_resizer = post_resizer
        self.device = device
        self.save_output = misc.SaveOutput()

        if self.eval_backbone == "InceptionV3_tf":
            self.res, mean, std = 299, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
            self.model = InceptionV3(resize_input=False, normalize_input=False).to(self.device)
        elif self.eval_backbone in ["InceptionV3_torch", "ResNet50_torch", "SwAV_torch"]:
            self.res = 299 if "InceptionV3" in self.eval_backbone else 224
            mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
            self.model = torch.hub.load(model_versions[self.eval_backbone],
                                        model_names[self.eval_backbone],
                                        pretrained=True)
            if self.eval_backbone == "SwAV_torch":
                linear_state_dict = load_state_dict_from_url(SWAV_CLASSIFIER_URL, progress=True)["state_dict"]
                linear_state_dict = {k.replace("module.linear.", ""): v for k, v in linear_state_dict.items()}
                self.model.fc.load_state_dict(linear_state_dict, strict=True)
            self.model = self.model.to(self.device)
            hook_handles = []
            for name, layer in self.model.named_children():
                if name == "fc":
                    handle = layer.register_forward_pre_hook(self.save_output)
                    hook_handles.append(handle)
        elif self.eval_backbone == "DINO_torch":
            self.res, mean, std = 224, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
            self.model = vits.__dict__["vit_small"](patch_size=8, num_classes=1000, num_last_blocks=4)
            misc.load_pretrained_weights(self.model, "", "teacher", "vit_small", 8)
            misc.load_pretrained_linear_weights(self.model.linear, "vit_small", 8)
            self.model = self.model.to(self.device)
        elif self.eval_backbone == "Swin-T_torch":
            self.res, mean, std = 224, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
            self.model = SwinTransformer()
            model_state_dict = load_state_dict_from_url(SWIN_URL, progress=True)["model"]
            self.model.load_state_dict(model_state_dict, strict=True)
            self.model = self.model.to(self.device)
        else:
            raise NotImplementedError

        self.resizer = resize.build_resizer(resizer=self.post_resizer, backbone=self.eval_backbone, size=self.res)
        self.totensor = transforms.ToTensor()
        self.mean = torch.Tensor(mean).view(1, 3, 1, 1).to(self.device)
        self.std = torch.Tensor(std).view(1, 3, 1, 1).to(self.device)

        if world_size > 1 and distributed_data_parallel:
            misc.make_model_require_grad(self.model)
            self.model = DDP(self.model,
                             device_ids=[self.device],
                             broadcast_buffers=False if self.eval_backbone=="Swin-T_torch" else True)
        elif world_size > 1 and distributed_data_parallel is False:
            self.model = DataParallel(self.model, output_device=self.device)
        else:
            pass

    def eval(self):
        self.model.eval()

    def get_outputs(self, x, quantize=False):
        if quantize:
            x = ops.quantize_images(x)
        else:
            x = x.detach().cpu().numpy().astype(np.uint8)
        x = ops.resize_images(x, self.resizer, self.totensor, self.mean, self.std, device=self.device)

        if self.eval_backbone in ["InceptionV3_tf", "DINO_torch", "Swin-T_torch"]:
            repres, logits = self.model(x)
        elif self.eval_backbone in ["InceptionV3_torch", "ResNet50_torch", "SwAV_torch"]:
            logits = self.model(x)
            if len(self.save_output.outputs) > 1:
                repres = []
                for rank in range(len(self.save_output.outputs)):
                    repres.append(self.save_output.outputs[rank][0].detach().cpu())
                repres = torch.cat(repres, dim=0).to(self.device)
            else:
                repres = self.save_output.outputs[0][0].to(self.device)
            self.save_output.clear()
        return repres, logits


def prepare_moments(data_loader, eval_model, quantize, cfgs, logger, device):
    disable_tqdm = device != 0
    eval_model.eval()
    moment_dir = join(cfgs.RUN.save_dir, "moments")
    if not exists(moment_dir):
        os.makedirs(moment_dir)
    moment_path = join(moment_dir, cfgs.DATA.name + "_"  + str(cfgs.DATA.img_size) + "_"+ cfgs.RUN.pre_resizer + "_" + \
                       cfgs.RUN.ref_dataset + "_" + cfgs.RUN.post_resizer + "_" + cfgs.RUN.eval_backbone + "_moments.npz")

    is_file = os.path.isfile(moment_path)
    if is_file:
        mu = np.load(moment_path)["mu"]
        sigma = np.load(moment_path)["sigma"]
    else:
        if device == 0:
            logger.info("Calculate moments of {ref} dataset using {eval_backbone} model.".\
                        format(ref=cfgs.RUN.ref_dataset, eval_backbone=cfgs.RUN.eval_backbone))
        mu, sigma = fid.calculate_moments(data_loader=data_loader,
                                          eval_model=eval_model,
                                          num_generate="N/A",
                                          batch_size=cfgs.OPTIMIZATION.batch_size,
                                          quantize=quantize,
                                          world_size=cfgs.OPTIMIZATION.world_size,
                                          DDP=cfgs.RUN.distributed_data_parallel,
                                          disable_tqdm=disable_tqdm,
                                          fake_feats=None)

        if device == 0:
            logger.info("Save calculated means and covariances to disk.")
        np.savez(moment_path, **{"mu": mu, "sigma": sigma})
    return mu, sigma


def prepare_real_feats(data_loader, eval_model, num_feats, quantize, cfgs, logger, device):
    disable_tqdm = device != 0
    eval_model.eval()
    feat_dir = join(cfgs.RUN.save_dir, "feats")
    if not exists(feat_dir):
        os.makedirs(feat_dir)
    feat_path = join(feat_dir, cfgs.DATA.name + "_"  + str(cfgs.DATA.img_size) + "_"+ cfgs.RUN.pre_resizer + "_" + \
                     cfgs.RUN.ref_dataset + "_" + cfgs.RUN.post_resizer + "_" + cfgs.RUN.eval_backbone + "_feats.npz")

    is_file = os.path.isfile(feat_path)
    if is_file:
        real_feats = np.load(feat_path)["real_feats"]
    else:
        if device == 0:
            logger.info("Calculate features of {ref} dataset using {eval_backbone} model.".\
                        format(ref=cfgs.RUN.ref_dataset, eval_backbone=cfgs.RUN.eval_backbone))
        real_feats, real_probs, real_labels = features.stack_features(data_loader=data_loader,
                                                eval_model=eval_model,
                                                num_feats=num_feats,
                                                batch_size=cfgs.OPTIMIZATION.batch_size,
                                                quantize=quantize,
                                                world_size=cfgs.OPTIMIZATION.world_size,
                                                DDP=cfgs.RUN.distributed_data_parallel,
                                                device=device,
                                                disable_tqdm=disable_tqdm)
        if device == 0:
            logger.info("Save real_features to disk.")
            np.savez(feat_path, **{"real_feats": real_feats,
                                   "real_probs": real_probs,
                                   "real_labels": real_labels})
    return real_feats


def calculate_ins(data_loader, eval_model, quantize, splits, cfgs, logger, device):
    disable_tqdm = device != 0
    is_acc = True if "ImageNet" in cfgs.DATA.name and "Tiny" not in cfgs.DATA.name else False
    if device == 0:
        logger.info("Calculate inception score of the {ref} dataset uisng pre-trained {eval_backbone} model.".\
                    format(ref=cfgs.RUN.ref_dataset, eval_backbone=cfgs.RUN.eval_backbone))
    is_score, is_std, top1, top5 = ins.eval_dataset(data_loader=data_loader,
                                                    eval_model=eval_model,
                                                    quantize=quantize,
                                                    splits=splits,
                                                    batch_size=cfgs.OPTIMIZATION.batch_size,
                                                    world_size=cfgs.OPTIMIZATION.world_size,
                                                    DDP=cfgs.RUN.distributed_data_parallel,
                                                    is_acc=is_acc,
                                                    is_torch_backbone=True if "torch" in cfgs.RUN.eval_backbone else False,
                                                    disable_tqdm=disable_tqdm)
    if device == 0:
        logger.info("Inception score={is_score}-Inception_std={is_std}".format(is_score=is_score, is_std=is_std))
        if is_acc:
            logger.info("{eval_model} Top1 acc: ({num} images): {Top1}".format(
                eval_model=cfgs.RUN.eval_backbone, num=str(len(data_loader.dataset)), Top1=top1))
            logger.info("{eval_model} Top5 acc: ({num} images): {Top5}".format(
                eval_model=cfgs.RUN.eval_backbone, num=str(len(data_loader.dataset)), Top5=top5))