+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 k 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.float().sum() / input.numel()
+ acc_sq += input.float().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 k in range(0, self.nb_batches):
+ self.inputs[k].sub_(mean).div_(std)
+
+ def get_batch(self, b):
+ return self.inputs[b], self.targets[b]
+
+class CompressedVignetteSet:
+ 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.input_storages = []
+
+ acc = 0.0
+ acc_sq = 0.0
+ for k in range(0, self.nb_batches):
+ target = torch.LongTensor(self.batch_size).bernoulli_(0.5)
+ input = svrt.generate_vignettes(problem_number, target)
+ acc += input.float().sum() / input.numel()
+ acc_sq += input.float().pow(2).sum() / input.numel()
+ self.targets.append(target)
+ self.input_storages.append(svrt.compress(input.storage()))
+
+ self.mean = acc / self.nb_batches
+ self.std = math.sqrt(acc_sq / self.nb_batches - self.mean * self.mean)
+
+ def get_batch(self, b):
+ input = torch.ByteTensor(svrt.uncompress(self.input_storages[b])).float()
+ input = input.view(self.batch_size, 1, 128, 128).sub_(self.mean).div_(self.std)
+ target = self.targets[b]
+
+ if torch.cuda.is_available():
+ input = input.cuda()
+ target = target.cuda()
+
+ return input, target