| import argparse |
| import pickle |
|
|
| import torch |
| from torch import nn |
| import numpy as np |
| from scipy import linalg |
| from tqdm import tqdm |
|
|
| from torchvision import transforms |
| from torchvision.datasets import ImageFolder |
| from torch.utils.data import DataLoader |
|
|
| from calc_inception import load_patched_inception_v3 |
| import os |
|
|
| @torch.no_grad() |
| def extract_features(loader, inception, device): |
| pbar = tqdm(loader) |
|
|
| feature_list = [] |
|
|
| for img,_ in pbar: |
| img = img.to(device) |
| feature = inception(img)[0].view(img.shape[0], -1) |
| feature_list.append(feature.to('cpu')) |
|
|
| features = torch.cat(feature_list, 0) |
|
|
| return features |
|
|
|
|
| def calc_fid(sample_mean, sample_cov, real_mean, real_cov, eps=1e-6): |
| cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False) |
|
|
| if not np.isfinite(cov_sqrt).all(): |
| print('product of cov matrices is singular') |
| offset = np.eye(sample_cov.shape[0]) * eps |
| cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset)) |
|
|
| if np.iscomplexobj(cov_sqrt): |
| if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3): |
| m = np.max(np.abs(cov_sqrt.imag)) |
|
|
| raise ValueError(f'Imaginary component {m}') |
|
|
| cov_sqrt = cov_sqrt.real |
|
|
| mean_diff = sample_mean - real_mean |
| mean_norm = mean_diff @ mean_diff |
|
|
| trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt) |
|
|
| fid = mean_norm + trace |
|
|
| return fid |
|
|
|
|
| if __name__ == '__main__': |
| device = 'cuda' |
|
|
| parser = argparse.ArgumentParser() |
|
|
| parser.add_argument('--batch', type=int, default=64) |
| parser.add_argument('--size', type=int, default=256) |
| parser.add_argument('--path_a', type=str) |
| parser.add_argument('--path_b', type=str) |
| parser.add_argument('--iter', type=int, default=3) |
| parser.add_argument('--end', type=int, default=13) |
|
|
| args = parser.parse_args() |
|
|
| inception = load_patched_inception_v3().eval().to(device) |
|
|
| transform = transforms.Compose( |
| [ |
| transforms.Resize( (args.size, args.size) ), |
| |
| transforms.ToTensor(), |
| transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), |
| ] |
| ) |
|
|
| dset_a = ImageFolder(args.path_a, transform) |
| loader_a = DataLoader(dset_a, batch_size=args.batch, num_workers=4) |
|
|
| features_a = extract_features(loader_a, inception, device).numpy() |
| print(f'extracted {features_a.shape[0]} features') |
|
|
| real_mean = np.mean(features_a, 0) |
| real_cov = np.cov(features_a, rowvar=False) |
| |
| |
| for folder in range(args.iter,args.end+1): |
| folder = 'eval_%d'%(folder*10000) |
| if os.path.exists(os.path.join( args.path_b, folder )): |
| print(folder) |
| dset_b = ImageFolder( os.path.join( args.path_b, folder ), transform) |
| loader_b = DataLoader(dset_b, batch_size=args.batch, num_workers=4) |
|
|
| features_b = extract_features(loader_b, inception, device).numpy() |
| print(f'extracted {features_b.shape[0]} features') |
|
|
| sample_mean = np.mean(features_b, 0) |
| sample_cov = np.cov(features_b, rowvar=False) |
|
|
| fid = calc_fid(sample_mean, sample_cov, real_mean, real_cov) |
|
|
| print(folder, ' fid:', fid) |
|
|
