[trainer] corrected stuff in the CNN trainer (saving models)

bob/learn/pytorch/trainers/CNNTrainer.py

@@ -13,7 +13,7 @@ import bob.core
 logger = bob.core.log.setup("bob.learn.pytorch")
 
 import time
-
+import os
 
 class CNNTrainer(object):
   """
@@ -75,7 +75,7 @@ class CNNTrainer(object):
     
     return start_epoch, start_iter, losses
 
-  def save_model(output_dir, epoch=0, iteration=0, losses=None):
+  def save_model(self, output_dir, epoch=0, iteration=0, losses=None):
     """
     Save the trained network
 
@@ -96,7 +96,7 @@ class CNNTrainer(object):
     saved_filename = 'model_{}_{}.pth'.format(epoch, iteration)    
     saved_path = os.path.join(output_dir, saved_filename)    
     
-    logger.info('Saving model to {}'.format(save_path))
+    logger.info('Saving model to {}'.format(saved_path))
     
     cp = {'epoch': epoch, 
           'iteration': iteration,
@@ -131,16 +131,19 @@ class CNNTrainer(object):
     # if model exists, load it
     if model is not None:
       start_epoch, start_iter, losses = self.load_model(model)
+      logger.info('Starting training at epoch {}, iteration {} - last loss value is {}'.format(start_epoch, start_iter, losses[-1]))
     else:
       start_epoch = 0
       start_iter = 0
       losses = []
+      logger.info('Starting training from scratch')
+
 
     # setup optimizer
     optimizer = optim.SGD(self.network.parameters(), learning_rate, momentum = 0.9, weight_decay = 0.0005)
 
     # let's go
-    for epoch in range((start_epoch, n_epochs):
+    for epoch in range(start_epoch, n_epochs):
       for i, data in enumerate(dataloader, 0):
    
         if i >= start_iter: