from .inception_v3 import InceptionV3 import blobfile as bf import torch import torch.distributed as dist import torch.nn as nn from cm import dist_util import numpy as np import warnings from scipy import linalg from PIL import Image from tqdm import tqdm def clip_preproc(preproc_fn, x): return preproc_fn(Image.fromarray(x.astype(np.uint8))) def all_gather(x, dim=0): xs = [torch.zeros_like(x) for _ in range(dist.get_world_size())] dist.all_gather(xs, x) return torch.cat(xs, dim=dim) class FIDStatistics: def __init__(self, mu: np.ndarray, sigma: np.ndarray, resolution: int): self.mu = mu self.sigma = sigma self.resolution = resolution def frechet_distance(self, other, eps=1e-6): """ Compute the Frechet distance between two sets of statistics. """ mu1, sigma1 = self.mu, self.sigma mu2, sigma2 = other.mu, other.sigma mu1 = np.atleast_1d(mu1) mu2 = np.atleast_1d(mu2) sigma1 = np.atleast_2d(sigma1) sigma2 = np.atleast_2d(sigma2) assert ( mu1.shape == mu2.shape ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}" assert ( sigma1.shape == sigma2.shape ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}" diff = mu1 - mu2 # product might be almost singular covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) if not np.isfinite(covmean).all(): msg = ( "fid calculation produces singular product; adding %s to diagonal of cov estimates" % eps ) warnings.warn(msg) offset = np.eye(sigma1.shape[0]) * eps covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) # numerical error might give slight imaginary component if np.iscomplexobj(covmean): if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): m = np.max(np.abs(covmean.imag)) raise ValueError("Imaginary component {}".format(m)) covmean = covmean.real tr_covmean = np.trace(covmean) return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean class FIDAndIS: def __init__( self, softmax_batch_size=512, clip_score_batch_size=512, path="https://openaipublic.blob.core.windows.net/consistency/inception/inception-2015-12-05.pt", ): import clip super().__init__() self.softmax_batch_size = softmax_batch_size self.clip_score_batch_size = clip_score_batch_size self.inception = InceptionV3() with bf.BlobFile(path, "rb") as f: self.inception.load_state_dict(torch.load(f)) self.inception.eval() self.inception.to(dist_util.dev()) self.inception_softmax = self.inception.create_softmax_model() if dist.get_rank() % 8 == 0: clip_model, self.clip_preproc_fn = clip.load( "ViT-B/32", device=dist_util.dev() ) dist.barrier() if dist.get_rank() % 8 != 0: clip_model, self.clip_preproc_fn = clip.load( "ViT-B/32", device=dist_util.dev() ) dist.barrier() # Compute the probe features separately from the final projection. class ProjLayer(nn.Module): def __init__(self, param): super().__init__() self.param = param def forward(self, x): return x @ self.param self.clip_visual = clip_model.visual self.clip_proj = ProjLayer(self.clip_visual.proj) self.clip_visual.proj = None class TextModel(nn.Module): def __init__(self, clip_model): super().__init__() self.clip_model = clip_model def forward(self, x): return self.clip_model.encode_text(x) self.clip_tokenizer = lambda captions: clip.tokenize(captions, truncate=True) self.clip_text = TextModel(clip_model) self.clip_logit_scale = clip_model.logit_scale.exp().item() self.ref_features = {} self.is_root = not dist.is_initialized() or dist.get_rank() == 0 def get_statistics(self, activations: np.ndarray, resolution: int): """ Compute activation statistics for a batch of images. :param activations: an [N x D] batch of activations. :return: an FIDStatistics object. """ mu = np.mean(activations, axis=0) sigma = np.cov(activations, rowvar=False) return FIDStatistics(mu, sigma, resolution) def get_preds(self, batch, captions=None): with torch.no_grad(): batch = 127.5 * (batch + 1) np_batch = batch.to(torch.uint8).cpu().numpy().transpose((0, 2, 3, 1)) pred, spatial_pred = self.inception(batch) pred, spatial_pred = pred.reshape( [pred.shape[0], -1] ), spatial_pred.reshape([spatial_pred.shape[0], -1]) clip_in = torch.stack( [clip_preproc(self.clip_preproc_fn, img) for img in np_batch] ) clip_pred = self.clip_visual(clip_in.half().to(dist_util.dev())) if captions is not None: text_in = self.clip_tokenizer(captions) text_pred = self.clip_text(text_in.to(dist_util.dev())) else: # Hack to easily deal with no captions text_pred = self.clip_proj(clip_pred.half()) text_pred = text_pred / text_pred.norm(dim=-1, keepdim=True) return pred, spatial_pred, clip_pred, text_pred, np_batch def get_inception_score( self, activations: np.ndarray, split_size: int = 5000 ) -> float: """ Compute the inception score using a batch of activations. :param activations: an [N x D] batch of activations. :param split_size: the number of samples per split. This is used to make results consistent with other work, even when using a different number of samples. :return: an inception score estimate. """ softmax_out = [] for i in range(0, len(activations), self.softmax_batch_size): acts = activations[i : i + self.softmax_batch_size] with torch.no_grad(): softmax_out.append( self.inception_softmax(torch.from_numpy(acts).to(dist_util.dev())) .cpu() .numpy() ) preds = np.concatenate(softmax_out, axis=0) # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46 scores = [] for i in range(0, len(preds), split_size): part = preds[i : i + split_size] kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0))) kl = np.mean(np.sum(kl, 1)) scores.append(np.exp(kl)) return float(np.mean(scores)) def get_clip_score( self, activations: np.ndarray, text_features: np.ndarray ) -> float: # Sizes should never mismatch, but if they do we want to compute # _some_ value instead of crash looping. size = min(len(activations), len(text_features)) activations = activations[:size] text_features = text_features[:size] scores_out = [] for i in range(0, len(activations), self.clip_score_batch_size): acts = activations[i : i + self.clip_score_batch_size] sub_features = text_features[i : i + self.clip_score_batch_size] with torch.no_grad(): image_features = self.clip_proj( torch.from_numpy(acts).half().to(dist_util.dev()) ) image_features = image_features / image_features.norm( dim=-1, keepdim=True ) image_features = image_features.detach().cpu().float().numpy() scores_out.extend(np.sum(sub_features * image_features, axis=-1).tolist()) return np.mean(scores_out) * self.clip_logit_scale def get_activations(self, data, num_samples, global_batch_size, pr_samples=50000): if self.is_root: preds = [] spatial_preds = [] clip_preds = [] pr_images = [] for _ in tqdm(range(0, int(np.ceil(num_samples / global_batch_size)))): batch, cond, _ = next(data) batch, cond = batch.to(dist_util.dev()), { k: v.to(dist_util.dev()) for k, v in cond.items() } pred, spatial_pred, clip_pred, _, np_batch = self.get_preds(batch) pred, spatial_pred, clip_pred = ( all_gather(pred).cpu().numpy(), all_gather(spatial_pred).cpu().numpy(), all_gather(clip_pred).cpu().numpy(), ) if self.is_root: preds.append(pred) spatial_preds.append(spatial_pred) clip_preds.append(clip_pred) if len(pr_images) * np_batch.shape[0] < pr_samples: pr_images.append(np_batch) if self.is_root: preds, spatial_preds, clip_preds, pr_images = ( np.concatenate(preds, axis=0), np.concatenate(spatial_preds, axis=0), np.concatenate(clip_preds, axis=0), np.concatenate(pr_images, axis=0), ) # assert len(pr_images) >= pr_samples return ( preds[:num_samples], spatial_preds[:num_samples], clip_preds[:num_samples], pr_images[:pr_samples], ) else: return [], [], [], [] def get_virtual_batch(self, data, num_samples, global_batch_size, resolution): preds, spatial_preds, clip_preds, batch = self.get_activations( data, num_samples, global_batch_size, pr_samples=10000 ) if self.is_root: fid_stats = self.get_statistics(preds, resolution) spatial_stats = self.get_statistics(spatial_preds, resolution) clip_stats = self.get_statistics(clip_preds, resolution) return batch, dict( mu=fid_stats.mu, sigma=fid_stats.sigma, mu_s=spatial_stats.mu, sigma_s=spatial_stats.sigma, mu_clip=clip_stats.mu, sigma_clip=clip_stats.sigma, ) else: return None, dict() def set_ref_batch(self, ref_batch): with bf.BlobFile(ref_batch, "rb") as f: data = np.load(f) fid_stats = FIDStatistics(mu=data["mu"], sigma=data["sigma"], resolution=-1) spatial_stats = FIDStatistics( mu=data["mu_s"], sigma=data["sigma_s"], resolution=-1 ) clip_stats = FIDStatistics( mu=data["mu_clip"], sigma=data["sigma_clip"], resolution=-1 ) self.ref_features[ref_batch] = (fid_stats, spatial_stats, clip_stats) def get_ref_batch(self, ref_batch): return self.ref_features[ref_batch]