X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=cnn-svrt.py;h=8baaacbc4fe7b4b37d8e6be27dd229b5b44bf6cc;hb=349b55a2d9ca213718df8941058d42689ba68163;hp=a41d42c41869009588f2c3e4d977503e0220e3bd;hpb=ce165da28eab42885032aa07e5defc6c72763576;p=pysvrt.git diff --git a/cnn-svrt.py b/cnn-svrt.py index a41d42c..8baaacb 100755 --- a/cnn-svrt.py +++ b/cnn-svrt.py @@ -32,12 +32,15 @@ from colorama import Fore, Back, Style # Pytorch import torch +import torchvision from torch import optim +from torch import multiprocessing from torch import FloatTensor as Tensor from torch.autograd import Variable from torch import nn from torch.nn import functional as fn + from torchvision import datasets, transforms, utils # SVRT @@ -73,13 +76,16 @@ parser.add_argument('--batch_size', parser.add_argument('--log_file', type = str, default = 'default.log') +parser.add_argument('--nb_exemplar_vignettes', + type = int, default = 32) + parser.add_argument('--compress_vignettes', type = distutils.util.strtobool, default = 'True', help = 'Use lossless compression to reduce the memory footprint') -parser.add_argument('--deep_model', - type = distutils.util.strtobool, default = 'True', - help = 'Use Afroze\'s Alexnet-like deep model') +parser.add_argument('--model', + type = str, default = 'deepnet', + help = 'What model to use') parser.add_argument('--test_loaded_models', type = distutils.util.strtobool, default = 'False', @@ -140,6 +146,8 @@ def log_string(s, remark = ''): # -- full(84x2) -> 2 1 class AfrozeShallowNet(nn.Module): + name = 'shallownet' + def __init__(self): super(AfrozeShallowNet, self).__init__() self.conv1 = nn.Conv2d(1, 6, kernel_size=21) @@ -147,7 +155,6 @@ class AfrozeShallowNet(nn.Module): 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)) @@ -163,6 +170,9 @@ class AfrozeShallowNet(nn.Module): # Afroze's DeepNet class AfrozeDeepNet(nn.Module): + + name = 'deepnet' + def __init__(self): super(AfrozeDeepNet, self).__init__() self.conv1 = nn.Conv2d( 1, 32, kernel_size=7, stride=4, padding=3) @@ -173,7 +183,6 @@ class AfrozeDeepNet(nn.Module): self.fc1 = nn.Linear(1536, 256) self.fc2 = nn.Linear(256, 256) self.fc3 = nn.Linear(256, 2) - self.name = 'deepnet' def forward(self, x): x = self.conv1(x) @@ -208,6 +217,108 @@ class AfrozeDeepNet(nn.Module): ###################################################################### +class DeepNet2(nn.Module): + name = 'deepnet2' + + def __init__(self): + super(DeepNet2, self).__init__() + self.conv1 = nn.Conv2d( 1, 32, kernel_size=7, stride=4, padding=3) + self.conv2 = nn.Conv2d( 32, 256, kernel_size=5, padding=2) + self.conv3 = nn.Conv2d(256, 256, kernel_size=3, padding=1) + self.conv4 = nn.Conv2d(256, 256, kernel_size=3, padding=1) + self.conv5 = nn.Conv2d(256, 256, kernel_size=3, padding=1) + self.fc1 = nn.Linear(4096, 512) + self.fc2 = nn.Linear(512, 512) + self.fc3 = nn.Linear(512, 2) + + def forward(self, x): + x = self.conv1(x) + x = fn.max_pool2d(x, kernel_size=2) + x = fn.relu(x) + + x = self.conv2(x) + x = fn.max_pool2d(x, kernel_size=2) + x = fn.relu(x) + + x = self.conv3(x) + x = fn.relu(x) + + x = self.conv4(x) + x = fn.relu(x) + + x = self.conv5(x) + x = fn.max_pool2d(x, kernel_size=2) + x = fn.relu(x) + + x = x.view(-1, 4096) + + x = self.fc1(x) + x = fn.relu(x) + + x = self.fc2(x) + x = fn.relu(x) + + x = self.fc3(x) + + return x + +###################################################################### + +class DeepNet3(nn.Module): + name = 'deepnet3' + + def __init__(self): + super(DeepNet3, self).__init__() + self.conv1 = nn.Conv2d( 1, 32, kernel_size=7, stride=4, padding=3) + self.conv2 = nn.Conv2d( 32, 128, kernel_size=5, padding=2) + self.conv3 = nn.Conv2d(128, 128, kernel_size=3, padding=1) + self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1) + self.conv5 = nn.Conv2d(128, 128, kernel_size=3, padding=1) + self.conv6 = nn.Conv2d(128, 128, kernel_size=3, padding=1) + self.conv7 = nn.Conv2d(128, 128, kernel_size=3, padding=1) + self.fc1 = nn.Linear(2048, 256) + self.fc2 = nn.Linear(256, 256) + self.fc3 = nn.Linear(256, 2) + + def forward(self, x): + x = self.conv1(x) + x = fn.max_pool2d(x, kernel_size=2) + x = fn.relu(x) + + x = self.conv2(x) + x = fn.max_pool2d(x, kernel_size=2) + x = fn.relu(x) + + x = self.conv3(x) + x = fn.relu(x) + + x = self.conv4(x) + x = fn.relu(x) + + x = self.conv5(x) + x = fn.max_pool2d(x, kernel_size=2) + x = fn.relu(x) + + x = self.conv6(x) + x = fn.relu(x) + + x = self.conv7(x) + x = fn.relu(x) + + x = x.view(-1, 2048) + + x = self.fc1(x) + x = fn.relu(x) + + x = self.fc2(x) + x = fn.relu(x) + + x = self.fc3(x) + + return x + +###################################################################### + def nb_errors(model, data_set): ne = 0 for b in range(0, data_set.nb_batches): @@ -223,7 +334,7 @@ def nb_errors(model, data_set): ###################################################################### -def train_model(model, train_set, validation_set): +def train_model(model, model_filename, train_set, validation_set, nb_epochs_done = 0): batch_size = args.batch_size criterion = nn.CrossEntropyLoss() @@ -234,7 +345,7 @@ def train_model(model, train_set, validation_set): start_t = time.time() - for e in range(0, args.nb_epochs): + for e in range(nb_epochs_done, args.nb_epochs): acc_loss = 0.0 for b in range(0, train_set.nb_batches): input, target = train_set.get_batch(b) @@ -249,6 +360,8 @@ def train_model(model, train_set, validation_set): log_string('train_loss {:d} {:f}'.format(e + 1, acc_loss), ' [ETA ' + time.ctime(time.time() + dt * (args.nb_epochs - e)) + ']') + torch.save([ model.state_dict(), e + 1 ], model_filename) + if validation_set is not None: nb_validation_errors = nb_errors(model, validation_set) @@ -295,6 +408,21 @@ class vignette_logger(): ) self.last_t = t +def save_examplar_vignettes(data_set, nb, name): + n = torch.randperm(data_set.nb_samples).narrow(0, 0, nb) + + for k in range(0, nb): + b = n[k] // data_set.batch_size + m = n[k] % data_set.batch_size + i, t = data_set.get_batch(b) + i = i[m].float() + i.sub_(i.min()) + i.div_(i.max()) + if k == 0: patchwork = Tensor(nb, 1, i.size(1), i.size(2)) + patchwork[k].copy_(i) + + torchvision.utils.save_image(patchwork, name) + ###################################################################### if args.nb_train_samples%args.batch_size > 0 or args.nb_test_samples%args.batch_size > 0: @@ -310,20 +438,30 @@ else: log_string('using_uncompressed_vignettes') VignetteSet = svrtset.VignetteSet +######################################## +model_class = None +for m in [ AfrozeShallowNet, AfrozeDeepNet, DeepNet2, DeepNet3 ]: + if args.model == m.name: + model_class = m + break +if model_class is None: + print('Unknown model ' + args.model) + raise + +log_string('using model class ' + m.name) +######################################## + for problem_number in map(int, args.problems.split(',')): log_string('############### problem ' + str(problem_number) + ' ###############') - if args.deep_model: - model = AfrozeDeepNet() - else: - model = AfrozeShallowNet() + model = model_class() if torch.cuda.is_available(): model.cuda() model_filename = model.name + '_pb:' + \ str(problem_number) + '_ns:' + \ - int_to_suffix(args.nb_train_samples) + '.param' + int_to_suffix(args.nb_train_samples) + '.state' nb_parameters = 0 for p in model.parameters(): nb_parameters += p.numel() @@ -332,17 +470,18 @@ for problem_number in map(int, args.problems.split(',')): ################################################## # Tries to load the model - need_to_train = False try: - model.load_state_dict(torch.load(model_filename)) + model_state_dict, nb_epochs_done = torch.load(model_filename) + model.load_state_dict(model_state_dict) log_string('loaded_model ' + model_filename) except: - need_to_train = True + nb_epochs_done = 0 + ################################################## # Train if necessary - if need_to_train: + if nb_epochs_done < args.nb_epochs: log_string('training_model ' + model_filename) @@ -357,6 +496,10 @@ for problem_number in map(int, args.problems.split(',')): train_set.nb_samples / (time.time() - t)) ) + if args.nb_exemplar_vignettes > 0: + save_examplar_vignettes(train_set, args.nb_exemplar_vignettes, + 'examplar_{:d}.png'.format(problem_number)) + if args.validation_error_threshold > 0.0: validation_set = VignetteSet(problem_number, args.nb_validation_samples, args.batch_size, @@ -365,8 +508,7 @@ for problem_number in map(int, args.problems.split(',')): else: validation_set = None - train_model(model, train_set, validation_set) - torch.save(model.state_dict(), model_filename) + train_model(model, model_filename, train_set, validation_set, nb_epochs_done = nb_epochs_done) log_string('saved_model ' + model_filename) nb_train_errors = nb_errors(model, train_set) @@ -381,7 +523,7 @@ for problem_number in map(int, args.problems.split(',')): ################################################## # Test if necessary - if need_to_train or args.test_loaded_models: + if nb_epochs_done < args.nb_epochs or args.test_loaded_models: t = time.time()