File size: 3,209 Bytes
c59c099 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import numpy as np
import torch
import torch.nn as nn
from scipy import linalg
from tqdm import tqdm
from basicsr.archs.inception import InceptionV3
def load_patched_inception_v3(device='cuda', resize_input=True, normalize_input=False):
# we may not resize the input, but in [rosinality/stylegan2-pytorch] it
# does resize the input.
inception = InceptionV3([3], resize_input=resize_input, normalize_input=normalize_input)
inception = nn.DataParallel(inception).eval().to(device)
return inception
@torch.no_grad()
def extract_inception_features(data_generator, inception, len_generator=None, device='cuda'):
"""Extract inception features.
Args:
data_generator (generator): A data generator.
inception (nn.Module): Inception model.
len_generator (int): Length of the data_generator to show the
progressbar. Default: None.
device (str): Device. Default: cuda.
Returns:
Tensor: Extracted features.
"""
if len_generator is not None:
pbar = tqdm(total=len_generator, unit='batch', desc='Extract')
else:
pbar = None
features = []
for data in data_generator:
if pbar:
pbar.update(1)
data = data.to(device)
feature = inception(data)[0].view(data.shape[0], -1)
features.append(feature.to('cpu'))
if pbar:
pbar.close()
features = torch.cat(features, 0)
return features
def calculate_fid(mu1, sigma1, mu2, sigma2, eps=1e-6):
"""Numpy implementation of the Frechet Distance.
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) and X_2 ~ N(mu_2, C_2) is:
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
Stable version by Dougal J. Sutherland.
Args:
mu1 (np.array): The sample mean over activations.
sigma1 (np.array): The covariance matrix over activations for generated samples.
mu2 (np.array): The sample mean over activations, precalculated on an representative data set.
sigma2 (np.array): The covariance matrix over activations, precalculated on an representative data set.
Returns:
float: The Frechet Distance.
"""
assert mu1.shape == mu2.shape, 'Two mean vectors have different lengths'
assert sigma1.shape == sigma2.shape, ('Two covariances have different dimensions')
cov_sqrt, _ = linalg.sqrtm(sigma1 @ sigma2, disp=False)
# Product might be almost singular
if not np.isfinite(cov_sqrt).all():
print('Product of cov matrices is singular. Adding {eps} to diagonal of cov estimates')
offset = np.eye(sigma1.shape[0]) * eps
cov_sqrt = linalg.sqrtm((sigma1 + offset) @ (sigma2 + offset))
# Numerical error might give slight imaginary component
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 = mu1 - mu2
mean_norm = mean_diff @ mean_diff
trace = np.trace(sigma1) + np.trace(sigma2) - 2 * np.trace(cov_sqrt)
fid = mean_norm + trace
return fid
|