Initial upload of main files

Author: ajbrock

animal_hash.py

@@ -0,0 +1,84 @@
+# This script iterates over the dataset and calculates the moments of the 
+# activations of the Inception net (needed for FID), and also returns
+# the Inception Score of the training data.
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import utils
+import inception_utils
+from tqdm import tqdm, trange
+from argparse import ArgumentParser
+
+def prepare_parser():
+  usage = 'Calculate and store inception metrics.'
+  parser = ArgumentParser(description=usage)
+  parser.add_argument(
+    '--dataset', type=str, default='I128',
+    help='Which Dataset to train on, out of I128, C10, C100, MN10, MN40, STL10 (default: %(default)s)') 
+  parser.add_argument(
+    '--batch_size', type=int, default=64,
+    help='Default overall batchsize (default: %(default)s)')
+  parser.add_argument(
+    '--parallel', action='store_true', default=False,
+    help='Train with multiple GPUs (default: %(default)s)')
+  parser.add_argument(
+    '--augment', action='store_true', default=False,
+    help='Augment with random crops and flips (default: %(default)s)')
+  parser.add_argument(
+    '--hdf5', action='store_true', default=False,
+    help='Use the HDF5 version of the dataset (default: %(default)s)') 
+  parser.add_argument(
+    '--num_workers', type=int, default=8,
+    help='Number of dataloader workers (default: %(default)s)')
+  parser.add_argument(
+    '--shuffle', action='store_true', default=False,
+    help='Shuffle the data? (default: %(default)s)') 
+  parser.add_argument(
+    '--seed', type=int, default=0,
+    help='Random seed to use.')
+  return parser
+
+def run(config):
+  # Get loader
+  loaders = utils.get_data_loaders(**config)
+
+  # Load inception net
+  net = inception_utils.load_inception_net(parallel=config['parallel'])
+  pool, logits, labels = [], [], []
+  device = 'cuda'
+  for i, (x, y) in enumerate(tqdm(loaders[0])):
+    x = x.to(device)
+    with torch.no_grad():
+      pool_val, logits_val = net(x)
+      pool += [np.asarray(pool_val.cpu())]
+      logits += [np.asarray(F.softmax(logits_val, 1).cpu())]
+      labels += [np.asarray(y.cpu())]
+
+  pool, logits, labels = [np.concatenate(item, 0) for item in [pool, logits, labels]]
+  # uncomment to save pool, logits, and labels to disk
+  # print('Saving pool, logits, and labels to disk...')
+  # np.savez(config['dataset']+'_inception_activations.npz', {'pool': pool, 'logits': logits, 'labels': labels})
+  # Calculate inception metrics and report them
+  print('Calculating inception metrics...')
+  IS_mean, IS_std = inception_utils.calculate_inception_score(logits)
+  print('Training data from dataset %s has IS of %5.5f +/- %5.5f' % (config['dataset'], IS_mean, IS_std))
+  # Prepare mu and sigma, save to disk
+  print('Calculating means and covariances...')
+  mu, sigma = np.mean(pool, axis=0), np.cov(pool, rowvar=False)
+  print('Saving calculated means and covariances to disk...')
+  np.savez(config['dataset']+'_inception_moments.npz', **{'mu' : mu, 'sigma' : sigma})
+
+def main():
+  # parse command line    
+  parser = prepare_parser()
+  config = vars(parser.parse_args())
+  print(config)
+  run(config)
+  # run; replace this with a for loop to run multiple sequential jobs.
+  # train_test(**vars(args))
+
+
+if __name__ == '__main__':    
+    main()