import time
import argparse
+import math
+
from colorama import Fore, Back, Style
import torch
formatter_class = argparse.ArgumentDefaultsHelpFormatter
)
-parser.add_argument('--nb_train_samples',
- type = int, default = 100000,
+parser.add_argument('--nb_train_batches',
+ type = int, default = 1000,
help = 'How many samples for train')
-parser.add_argument('--nb_test_samples',
- type = int, default = 10000,
+parser.add_argument('--nb_test_batches',
+ type = int, default = 100,
help = 'How many samples for test')
parser.add_argument('--nb_epochs',
- type = int, default = 25,
+ 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()
######################################################################
print(Fore.RED + 'Logging into ' + args.log_file + Style.RESET_ALL)
def log_string(s):
- s = Fore.GREEN + time.ctime() + Style.RESET_ALL + ' ' + \
- str(problem_number) + ' ' + s
+ s = Fore.GREEN + time.ctime() + Style.RESET_ALL + ' ' + s
log_file.write(s + '\n')
log_file.flush()
print(s)
######################################################################
-def generate_set(p, n):
- target = torch.LongTensor(n).bernoulli_(0.5)
- t = time.time()
- input = svrt.generate_vignettes(p, target)
- t = time.time() - t
- log_string('DATA_SET_GENERATION {:.02f} sample/s'.format(n / t))
- input = input.view(input.size(0), 1, input.size(1), input.size(2)).float()
- return Variable(input), Variable(target)
+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
-# ----------------------
-# 128x128 1
-# -- conv(21x21) -> 108x108 6
-# -- max(2x2) -> 54x54 6
-# -- conv(19x19) -> 36x36 16
-# -- max(2x2) -> 18x18 16
-# -- conv(18x18) -> 1x1 120
-# -- reshape -> 120 1
-# -- full(120x84) -> 84 1
-# -- full(84x2) -> 2 1
-
-class Net(nn.Module):
+# map size nb. maps
+# ----------------------
+# input 128x128 1
+# -- conv(21x21 x 6) -> 108x108 6
+# -- max(2x2) -> 54x54 6
+# -- conv(19x19 x 16) -> 36x36 16
+# -- max(2x2) -> 18x18 16
+# -- conv(18x18 x 120) -> 1x1 120
+# -- reshape -> 120 1
+# -- full(120x84) -> 84 1
+# -- full(84x2) -> 2 1
+
+class AfrozeShallowNet(nn.Module):
def __init__(self):
- super(Net, self).__init__()
+ super(AfrozeShallowNet, self).__init__()
self.conv1 = nn.Conv2d(1, 6, kernel_size=21)
self.conv2 = nn.Conv2d(6, 16, kernel_size=19)
self.conv3 = nn.Conv2d(16, 120, kernel_size=18)
x = self.fc2(x)
return x
-def train_model(train_input, train_target):
- model, criterion = Net(), nn.CrossEntropyLoss()
+def train_model(model, train_set):
+ batch_size = args.batch_size
+ criterion = nn.CrossEntropyLoss()
if torch.cuda.is_available():
- model.cuda()
criterion.cuda()
- optimizer, bs = optim.SGD(model.parameters(), lr = 1e-2), 100
+ optimizer = optim.SGD(model.parameters(), lr = 1e-2)
- for k in range(0, args.nb_epochs):
+ for e in range(0, args.nb_epochs):
acc_loss = 0.0
- for b in range(0, train_input.size(0), bs):
- output = model.forward(train_input.narrow(0, b, bs))
- loss = criterion(output, train_target.narrow(0, b, bs))
+ for b in range(0, train_set.nb_batches):
+ input, target = train_set.get_batch(b)
+ output = model.forward(Variable(input))
+ loss = criterion(output, Variable(target))
acc_loss = acc_loss + loss.data[0]
model.zero_grad()
loss.backward()
optimizer.step()
- log_string('TRAIN_LOSS {:d} {:f}'.format(k, acc_loss))
+ log_string('train_loss {:d} {:f}'.format(e + 1, acc_loss))
return model
######################################################################
-def nb_errors(model, data_input, data_target, bs = 100):
+def nb_errors(model, data_set):
ne = 0
-
- for b in range(0, data_input.size(0), bs):
- output = model.forward(data_input.narrow(0, b, bs))
+ for b in range(0, data_set.nb_batches):
+ input, target = data_set.get_batch(b)
+ output = model.forward(Variable(input))
wta_prediction = output.data.max(1)[1].view(-1)
- for i in range(0, bs):
- if wta_prediction[i] != data_target.narrow(0, b, bs).data[i]:
+ for i in range(0, data_set.batch_size):
+ if wta_prediction[i] != target[i]:
ne = ne + 1
return ne
######################################################################
+for arg in vars(args):
+ log_string('argument ' + str(arg) + ' ' + str(getattr(args, arg)))
+
for problem_number in range(1, 24):
- train_input, train_target = generate_set(problem_number, args.nb_train_samples)
- test_input, test_target = generate_set(problem_number, args.nb_test_samples)
+ if args.compress_vignettes:
+ train_set = CompressedVignetteSet(problem_number, args.nb_train_batches)
+ test_set = CompressedVignetteSet(problem_number, args.nb_test_batches)
+ else:
+ train_set = VignetteSet(problem_number, args.nb_train_batches)
+ test_set = VignetteSet(problem_number, args.nb_test_batches)
+
+ model = AfrozeShallowNet()
if torch.cuda.is_available():
- train_input, train_target = train_input.cuda(), train_target.cuda()
- test_input, test_target = test_input.cuda(), test_target.cuda()
+ model.cuda()
+
+ nb_parameters = 0
+ for p in model.parameters():
+ nb_parameters += p.numel()
+ log_string('nb_parameters {:d}'.format(nb_parameters))
- mu, std = train_input.data.mean(), train_input.data.std()
- train_input.data.sub_(mu).div_(std)
- test_input.data.sub_(mu).div_(std)
+ model_filename = 'model_' + str(problem_number) + '.param'
- model = train_model(train_input, train_target)
+ try:
+ model.load_state_dict(torch.load(model_filename))
+ log_string('loaded_model ' + model_filename)
+ except:
+ log_string('training_model')
+ train_model(model, train_set)
+ torch.save(model.state_dict(), model_filename)
+ log_string('saved_model ' + model_filename)
- nb_train_errors = nb_errors(model, train_input, train_target)
+ nb_train_errors = nb_errors(model, train_set)
- log_string('TRAIN_ERROR {:.02f}% {:d} {:d}'.format(
- 100 * nb_train_errors / train_input.size(0),
+ log_string('train_error {:d} {:.02f}% {:d} {:d}'.format(
+ problem_number,
+ 100 * nb_train_errors / train_set.nb_samples,
nb_train_errors,
- train_input.size(0))
+ train_set.nb_samples)
)
- nb_test_errors = nb_errors(model, test_input, test_target)
+ nb_test_errors = nb_errors(model, test_set)
- log_string('TEST_ERROR {:.02f}% {:d} {:d}'.format(
- 100 * nb_test_errors / test_input.size(0),
+ log_string('test_error {:d} {:.02f}% {:d} {:d}'.format(
+ problem_number,
+ 100 * nb_test_errors / test_set.nb_samples,
nb_test_errors,
- test_input.size(0))
+ test_set.nb_samples)
)
######################################################################
projects / pysvrt.git / blobdiff