+
+######################################################################
+
+parser = argparse.ArgumentParser(
+ description = 'An implementation of Mutual Information estimator with a deep model',
+ formatter_class = argparse.ArgumentDefaultsHelpFormatter
+)
+
+parser.add_argument('--data',
+ type = str, default = 'image_pair',
+ help = 'What data')
+
+parser.add_argument('--seed',
+ type = int, default = 0,
+ help = 'Random seed (default 0, < 0 is no seeding)')
+
+parser.add_argument('--mnist_classes',
+ type = str, default = '0, 1, 3, 5, 6, 7, 8, 9',
+ help = 'What MNIST classes to use')
+
+######################################################################
+
+if torch.cuda.is_available():
+ device = torch.device('cuda')
+else:
+ device = torch.device('cpu')
+
+######################################################################
+
+def entropy(target):
+ probas = []
+ for k in range(target.max() + 1):
+ n = (target == k).sum().item()
+ if n > 0: probas.append(n)
+ probas = torch.tensor(probas).float()
+ probas /= probas.sum()
+ return - (probas * probas.log()).sum().item()
+
+######################################################################
+
+args = parser.parse_args()
+
+if args.seed >= 0:
+ torch.manual_seed(args.seed)
+
+used_MNIST_classes = torch.tensor(eval('[' + args.mnist_classes + ']'), device = device)