+parser = argparse.ArgumentParser(
+ description = 'Simple convnet test on the SVRT.',
+ formatter_class = argparse.ArgumentDefaultsHelpFormatter
+)
+
+parser.add_argument('--nb_train_batches',
+ type = int, default = 1000,
+ help = 'How many samples for train')
+
+parser.add_argument('--nb_test_batches',
+ type = int, default = 100,
+ help = 'How many samples for test')
+
+parser.add_argument('--nb_epochs',
+ type = int, default = 50,
+ help = 'How many training epochs')
+
+parser.add_argument('--batch_size',
+ type = int, default = 100,
+ help = 'Mini-batch size')
+
+parser.add_argument('--log_file',
+ type = str, default = 'cnn-svrt.log',
+ help = 'Log file name')
+
+parser.add_argument('--compress_vignettes',
+ action='store_true', default = False,
+ help = 'Use lossless compression to reduce the memory footprint')
+
+args = parser.parse_args()
+
+######################################################################
+
+log_file = open(args.log_file, 'w')
+
+print(Fore.RED + 'Logging into ' + args.log_file + Style.RESET_ALL)
+
+def log_string(s):
+ s = Fore.GREEN + time.ctime() + Style.RESET_ALL + ' ' + s
+ log_file.write(s + '\n')
+ log_file.flush()
+ print(s)
+
+######################################################################
+
+class VignetteSet:
+ def __init__(self, problem_number, nb_batches):
+ self.batch_size = args.batch_size
+ self.problem_number = problem_number
+ self.nb_batches = nb_batches
+ self.nb_samples = self.nb_batches * self.batch_size
+ self.targets = []
+ self.inputs = []
+
+ acc = 0.0
+ acc_sq = 0.0
+
+ for b in range(0, self.nb_batches):
+ target = torch.LongTensor(self.batch_size).bernoulli_(0.5)
+ input = svrt.generate_vignettes(problem_number, target)
+ input = input.float().view(input.size(0), 1, input.size(1), input.size(2))
+ if torch.cuda.is_available():
+ input = input.cuda()
+ target = target.cuda()
+ acc += input.sum() / input.numel()
+ acc_sq += input.pow(2).sum() / input.numel()
+ self.targets.append(target)
+ self.inputs.append(input)
+
+ mean = acc / self.nb_batches
+ std = math.sqrt(acc_sq / self.nb_batches - mean * mean)
+ for b in range(0, self.nb_batches):
+ self.inputs[b].sub_(mean).div_(std)
+
+ def get_batch(self, b):
+ return self.inputs[b], self.targets[b]