Minor update.

[pysvrt.git] / cnn-svrt.py

diff --git a/cnn-svrt.py b/cnn-svrt.py
index a41d42c..a6b9cab 100755 (executable)
--- a/cnn-svrt.py
+++ b/cnn-svrt.py
@@ -24,20 +24,25 @@
 import time
 import argparse
 import math
 import time
 import argparse
 import math

+

 import distutils.util
 import re
 import distutils.util
 import re

+import signal

 
 from colorama import Fore, Back, Style
 
 # Pytorch
 
 import torch
 
 from colorama import Fore, Back, Style
 
 # Pytorch
 
 import torch

+import torchvision

 
 from torch import optim
 
 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 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
 from torchvision import datasets, transforms, utils
 
 # SVRT


@@ -73,13 +78,19 @@ parser.add_argument('--batch_size',
 parser.add_argument('--log_file',
                     type = str, default = 'default.log')
 
 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('--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('--save_test_mistakes',
+                    type = distutils.util.strtobool, default = 'False')
+
+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',
 
 parser.add_argument('--test_loaded_models',
                     type = distutils.util.strtobool, default = 'False',


@@ -94,6 +105,10 @@ args = parser.parse_args()
 ######################################################################
 
 log_file = open(args.log_file, 'a')
 ######################################################################
 
 log_file = open(args.log_file, 'a')

+log_file.write('\n')
+log_file.write('@@@@@@@@@@@@@@@@@@@ ' + time.ctime() + ' @@@@@@@@@@@@@@@@@@@\n')
+log_file.write('\n')
+

 pred_log_t = None
 last_tag_t = time.time()
 
 pred_log_t = None
 last_tag_t = time.time()
 


@@ -121,7 +136,24 @@ def log_string(s, remark = ''):
     log_file.write(re.sub(' ', '_', time.ctime()) + ' ' + elapsed + ' ' + s + '\n')
     log_file.flush()
 
     log_file.write(re.sub(' ', '_', time.ctime()) + ' ' + elapsed + ' ' + s + '\n')
     log_file.flush()
 

-    print(Fore.BLUE + time.ctime() + ' ' + Fore.GREEN + elapsed + Style.RESET_ALL + ' ' + s + Fore.CYAN + remark + Style.RESET_ALL)
+    print(Fore.BLUE + time.ctime() + ' ' + Fore.GREEN + elapsed \
+          + Style.RESET_ALL
+          + ' ' \
+          + s + Fore.CYAN + remark \
+          + Style.RESET_ALL)
+
+######################################################################
+
+def handler_sigint(signum, frame):
+    log_string('got sigint')
+    exit(0)
+
+def handler_sigterm(signum, frame):
+    log_string('got sigterm')
+    exit(0)
+
+signal.signal(signal.SIGINT, handler_sigint)
+signal.signal(signal.SIGTERM, handler_sigterm)

 
 ######################################################################
 
 
 ######################################################################
 


@@ -140,6 +172,8 @@ def log_string(s, remark = ''):
 # -- full(84x2)        -> 2          1
 
 class AfrozeShallowNet(nn.Module):
 # -- 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)
     def __init__(self):
         super(AfrozeShallowNet, self).__init__()
         self.conv1 = nn.Conv2d(1, 6, kernel_size=21)


@@ -147,7 +181,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.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))
 
     def forward(self, x):
         x = fn.relu(fn.max_pool2d(self.conv1(x), kernel_size=2))


@@ -163,6 +196,9 @@ class AfrozeShallowNet(nn.Module):
 # Afroze's DeepNet
 
 class AfrozeDeepNet(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)
     def __init__(self):
         super(AfrozeDeepNet, self).__init__()
         self.conv1 = nn.Conv2d(  1,  32, kernel_size=7, stride=4, padding=3)


@@ -173,7 +209,6 @@ class AfrozeDeepNet(nn.Module):
         self.fc1 = nn.Linear(1536, 256)
         self.fc2 = nn.Linear(256, 256)
         self.fc3 = nn.Linear(256, 2)
         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)
 
     def forward(self, x):
         x = self.conv1(x)


@@ -208,7 +243,110 @@ class AfrozeDeepNet(nn.Module):
 
 ######################################################################
 
 
 ######################################################################
 

-def nb_errors(model, data_set):
+class DeepNet2(nn.Module):
+    name = 'deepnet2'
+
+    def __init__(self):
+        super(DeepNet2, self).__init__()
+        self.nb_channels = 512
+        self.conv1 = nn.Conv2d(  1,  32, kernel_size=7, stride=4, padding=3)
+        self.conv2 = nn.Conv2d( 32, self.nb_channels, kernel_size=5, padding=2)
+        self.conv3 = nn.Conv2d(self.nb_channels, self.nb_channels, kernel_size=3, padding=1)
+        self.conv4 = nn.Conv2d(self.nb_channels, self.nb_channels, kernel_size=3, padding=1)
+        self.conv5 = nn.Conv2d(self.nb_channels, self.nb_channels, kernel_size=3, padding=1)
+        self.fc1 = nn.Linear(16 * self.nb_channels, 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, 16 * self.nb_channels)
+
+        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, mistake_filename_pattern = None):

     ne = 0
     for b in range(0, data_set.nb_batches):
         input, target = data_set.get_batch(b)
     ne = 0
     for b in range(0, data_set.nb_batches):
         input, target = data_set.get_batch(b)


@@ -218,12 +356,19 @@ def nb_errors(model, data_set):
         for i in range(0, data_set.batch_size):
             if wta_prediction[i] != target[i]:
                 ne = ne + 1
         for i in range(0, data_set.batch_size):
             if wta_prediction[i] != target[i]:
                 ne = ne + 1

-
+                if mistake_filename_pattern is not None:
+                    img = input[i].clone()
+                    img.sub_(img.min())
+                    img.div_(img.max())
+                    k = b * data_set.batch_size + i
+                    filename = mistake_filename_pattern.format(k, target[i])
+                    torchvision.utils.save_image(img, filename)
+                    print(Fore.RED + 'Wrote ' + filename + Style.RESET_ALL)

     return ne
 
 ######################################################################
 
     return ne
 
 ######################################################################
 

-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()
 
     batch_size = args.batch_size
     criterion = nn.CrossEntropyLoss()
 


@@ -234,7 +379,7 @@ def train_model(model, train_set, validation_set):
 
     start_t = time.time()
 
 
     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)
         acc_loss = 0.0
         for b in range(0, train_set.nb_batches):
             input, target = train_set.get_batch(b)


@@ -249,6 +394,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)) + ']')
 
         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)
 
         if validation_set is not None:
             nb_validation_errors = nb_errors(model, validation_set)
 


@@ -295,14 +442,27 @@ class vignette_logger():
             )
             self.last_t = t
 
             )
             self.last_t = t
 

+def save_exemplar_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:
     print('The number of samples must be a multiple of the batch size.')
     raise
 
 ######################################################################
 
 if args.nb_train_samples%args.batch_size > 0 or args.nb_test_samples%args.batch_size > 0:
     print('The number of samples must be a multiple of the batch size.')
     raise
 

-log_string('############### start ###############')
-

 if args.compress_vignettes:
     log_string('using_compressed_vignettes')
     VignetteSet = svrtset.CompressedVignetteSet
 if args.compress_vignettes:
     log_string('using_compressed_vignettes')
     VignetteSet = svrtset.CompressedVignetteSet


@@ -310,20 +470,30 @@ else:
     log_string('using_uncompressed_vignettes')
     VignetteSet = svrtset.VignetteSet
 
     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) + ' ###############')
 
 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:' + \
 
     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) + '.pth'

 
     nb_parameters = 0
     for p in model.parameters(): nb_parameters += p.numel()
 
     nb_parameters = 0
     for p in model.parameters(): nb_parameters += p.numel()


@@ -332,17 +502,18 @@ for problem_number in map(int, args.problems.split(',')):
     ##################################################
     # Tries to load the model
 
     ##################################################
     # Tries to load the model
 

-    need_to_train = False

     try:
     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:
         log_string('loaded_model ' + model_filename)
     except:

-        need_to_train = True
+        nb_epochs_done = 0
+

 
     ##################################################
     # Train if necessary
 
 
     ##################################################
     # Train if necessary
 

-    if need_to_train:
+    if nb_epochs_done < args.nb_epochs:

 
         log_string('training_model ' + model_filename)
 
 
         log_string('training_model ' + model_filename)
 


@@ -357,6 +528,10 @@ for problem_number in map(int, args.problems.split(',')):
             train_set.nb_samples / (time.time() - t))
         )
 
             train_set.nb_samples / (time.time() - t))
         )
 

+        if args.nb_exemplar_vignettes > 0:
+            save_exemplar_vignettes(train_set, args.nb_exemplar_vignettes,
+                                    'exemplar_{:d}.png'.format(problem_number))
+

         if args.validation_error_threshold > 0.0:
             validation_set = VignetteSet(problem_number,
                                          args.nb_validation_samples, args.batch_size,
         if args.validation_error_threshold > 0.0:
             validation_set = VignetteSet(problem_number,
                                          args.nb_validation_samples, args.batch_size,


@@ -365,8 +540,10 @@ for problem_number in map(int, args.problems.split(',')):
         else:
             validation_set = None
 
         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)
         log_string('saved_model ' + model_filename)
 
         nb_train_errors = nb_errors(model, train_set)


@@ -381,7 +558,7 @@ for problem_number in map(int, args.problems.split(',')):
     ##################################################
     # Test if necessary
 
     ##################################################
     # 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()
 
 
         t = time.time()
 


@@ -389,7 +566,8 @@ for problem_number in map(int, args.problems.split(',')):
                                args.nb_test_samples, args.batch_size,
                                cuda = torch.cuda.is_available())
 
                                args.nb_test_samples, args.batch_size,
                                cuda = torch.cuda.is_available())
 

-        nb_test_errors = nb_errors(model, test_set)
+        nb_test_errors = nb_errors(model, test_set,
+                                   mistake_filename_pattern = 'mistake_{:06d}_{:d}.png')

 
         log_string('test_error {:d} {:.02f}% {:d} {:d}'.format(
             problem_number,
 
         log_string('test_error {:d} {:.02f}% {:d} {:d}'.format(
             problem_number,