| |
| |
| |
| |
| |
| |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from scipy import linalg |
| from tqdm import tqdm |
|
|
| from basicsr.models.archs.inception import InceptionV3 |
|
|
|
|
| def load_patched_inception_v3(device='cuda', |
| resize_input=True, |
| normalize_input=False): |
| |
| |
| 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) |
|
|
| |
| 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)) |
|
|
| |
| 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 |
|
|
