Heavy fix.
[pysvrt.git] / vignette_set.py
1
2 #  svrt is the ``Synthetic Visual Reasoning Test'', an image
3 #  generator for evaluating classification performance of machine
4 #  learning systems, humans and primates.
5 #
6 #  Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/
7 #  Written by Francois Fleuret <francois.fleuret@idiap.ch>
8 #
9 #  This file is part of svrt.
10 #
11 #  svrt is free software: you can redistribute it and/or modify it
12 #  under the terms of the GNU General Public License version 3 as
13 #  published by the Free Software Foundation.
14 #
15 #  svrt is distributed in the hope that it will be useful, but
16 #  WITHOUT ANY WARRANTY; without even the implied warranty of
17 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
18 #  General Public License for more details.
19 #
20 #  You should have received a copy of the GNU General Public License
21 #  along with selector.  If not, see <http://www.gnu.org/licenses/>.
22
23 import torch
24 from math import sqrt
25 from multiprocessing import Pool, cpu_count
26
27 from torch import Tensor
28 from torch.autograd import Variable
29
30 import svrt
31
32 ######################################################################
33
34 def generate_one_batch(s):
35     problem_number, batch_size, cuda, random_seed = s
36     svrt.seed(random_seed)
37     target = torch.LongTensor(batch_size).bernoulli_(0.5)
38     input = svrt.generate_vignettes(problem_number, target)
39     input = input.float().view(input.size(0), 1, input.size(1), input.size(2))
40     if cuda:
41         input = input.cuda()
42         target = target.cuda()
43     return [ input, target ]
44
45 class VignetteSet:
46
47     def __init__(self, problem_number, nb_batches, batch_size, cuda = False):
48         self.cuda = cuda
49         self.batch_size = batch_size
50         self.problem_number = problem_number
51         self.nb_batches = nb_batches
52         self.nb_samples = self.nb_batches * self.batch_size
53
54         seeds = torch.LongTensor(self.nb_batches).random_()
55         mp_args = []
56         for b in range(0, self.nb_batches):
57             mp_args.append( [ problem_number, batch_size, cuda, seeds[b] ])
58
59         # self.data = []
60         # for b in range(0, self.nb_batches):
61             # self.data.append(generate_one_batch(mp_args[b]))
62
63         self.data = Pool(cpu_count()).map(generate_one_batch, mp_args)
64
65         acc = 0.0
66         acc_sq = 0.0
67         for b in range(0, self.nb_batches):
68             input = self.data[b][0]
69             acc += input.sum() / input.numel()
70             acc_sq += input.pow(2).sum() /  input.numel()
71
72         mean = acc / self.nb_batches
73         std = sqrt(acc_sq / self.nb_batches - mean * mean)
74         for b in range(0, self.nb_batches):
75             self.data[b][0].sub_(mean).div_(std)
76
77     def get_batch(self, b):
78         return self.data[b]
79
80 ######################################################################
81
82 class CompressedVignetteSet:
83     def __init__(self, problem_number, nb_batches, batch_size, cuda = False):
84         self.cuda = cuda
85         self.batch_size = batch_size
86         self.problem_number = problem_number
87         self.nb_batches = nb_batches
88         self.nb_samples = self.nb_batches * self.batch_size
89         self.targets = []
90         self.input_storages = []
91
92         acc = 0.0
93         acc_sq = 0.0
94         for b in range(0, self.nb_batches):
95             target = torch.LongTensor(self.batch_size).bernoulli_(0.5)
96             input = svrt.generate_vignettes(problem_number, target)
97             acc += input.float().sum() / input.numel()
98             acc_sq += input.float().pow(2).sum() /  input.numel()
99             self.targets.append(target)
100             self.input_storages.append(svrt.compress(input.storage()))
101
102         self.mean = acc / self.nb_batches
103         self.std = sqrt(acc_sq / self.nb_batches - self.mean * self.mean)
104
105     def get_batch(self, b):
106         input = torch.ByteTensor(svrt.uncompress(self.input_storages[b])).float()
107         input = input.view(self.batch_size, 1, 128, 128).sub_(self.mean).div_(self.std)
108         target = self.targets[b]
109
110         if self.cuda:
111             input = input.cuda()
112             target = target.cuda()
113
114         return input, target
115
116 ######################################################################