text: import argparse
  import pickle
  
  import torch
  from torch import nn
  import numpy as np
  from scipy import linalg
  from tqdm import tqdm
  
  from model import Generator
  from calc_inception import load_patched_inception_v3
  
  
  @torch.no_grad()
  def extract_feature_from_samples(
      generator, inception, truncation, truncation_latent, batch_size, n_sample, device
  ):
      n_batch = n_sample // batch_size
      resid = n_sample - (n_batch * batch_size)
      batch_sizes = [batch_size] * n_batch + [resid]
      features = []
  
      for batch in tqdm(batch_sizes):
          latent = torch.randn(batch, 256, device=device)
          img, _ = g([latent], truncation=truncation, truncation_latent=truncation_latent)
          feat = inception(img)[0].view(img.shape[0], -1)
          features.append(feat.to('cpu'))
  
      features = torch.cat(features, 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('--truncation', type=float, default=1)
      parser.add_argument('--truncation_mean', type=int, default=4096)
      parser.add_argument('--batch', type=int, default=64)
      parser.add_argument('--n_sample', type=int, default=50000)
      parser.add_argument('--size', type=int, default=256)
      parser.add_argument('--inception', type=str, default=None, required=True)
      parser.add_argument('ckpt', metavar='CHECKPOINT')
  
      args = parser.parse_args()
  
      ckpt = torch.load(args.ckpt)
  
      g = Generator(args.size, 256, 4, channel_multiplier=1).to(device)
      g.load_state_dict(ckpt['g_ema'])
      g = nn.DataParallel(g)
      g.eval()
  
      if args.truncation < 1:
          with torch.no_grad():
              mean_latent = g.mean_latent(args.truncation_mean)
  
      else:
          mean_latent = None
  
      inception = nn.DataParallel(load_patched_inception_v3()).to(device)
      inception.eval()
  
      features = extract_feature_from_samples(
          g, inception, args.truncation, mean_latent, args.batch, args.n_sample, device
      ).numpy()
      print(f'extracted {features.shape[0]} features')
  
      sample_mean = np.mean(features, 0)
      sample_cov = np.cov(features, rowvar=False)
  
      with open(args.inception, 'rb') as f:
          embeds = pickle.load(f)
          real_mean = embeds['mean']
          real_cov = embeds['cov']
  
      fid = calc_fid(sample_mean, sample_cov, real_mean, real_cov)
  
      print('fid:', fid)
  