from torch.nn import functional as fn
from torchvision import datasets, transforms, utils
-import svrt
+from vignette_set import VignetteSet, CompressedVignetteSet
######################################################################
######################################################################
-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]
-
-######################################################################
-
-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 b 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
-
-######################################################################
-
# Afroze's ShallowNet
# map size nb. maps
self.conv3 = nn.Conv2d(16, 120, kernel_size=18)
self.fc1 = nn.Linear(120, 84)
self.fc2 = nn.Linear(84, 2)
+ self.name = 'shallownet'
def forward(self, x):
x = fn.relu(fn.max_pool2d(self.conv1(x), kernel_size=2))
x = self.fc2(x)
return x
+######################################################################
+
def train_model(model, train_set):
batch_size = args.batch_size
criterion = nn.CrossEntropyLoss()
for problem_number in range(1, 24):
if args.compress_vignettes:
- train_set = CompressedVignetteSet(problem_number, args.nb_train_batches)
- test_set = CompressedVignetteSet(problem_number, args.nb_test_batches)
+ train_set = CompressedVignetteSet(problem_number, args.nb_train_batches, args.batch_size)
+ test_set = CompressedVignetteSet(problem_number, args.nb_test_batches, args.batch_size)
else:
- train_set = VignetteSet(problem_number, args.nb_train_batches)
- test_set = VignetteSet(problem_number, args.nb_test_batches)
+ train_set = VignetteSet(problem_number, args.nb_train_batches, args.batch_size)
+ test_set = VignetteSet(problem_number, args.nb_test_batches, args.batch_size)
model = AfrozeShallowNet()
nb_parameters += p.numel()
log_string('nb_parameters {:d}'.format(nb_parameters))
- model_filename = 'model_' + str(problem_number) + '.param'
+ model_filename = model.name + '_' + str(problem_number) + '_' + str(train_set.nb_batches) + '.param'
try:
model.load_state_dict(torch.load(model_filename))