update training

parameters_combinations/param_combinations_chairs_test.txt

---batch-size=64 --dataset=rendered_chairs --epochs=4 --latent_spec_cont=10 --lr=1e-4 --experiment-name=VAE_bs_64_conv_64_64_128_128 --gpu-devices 0 1 --nb-filter-conv1=64 --nb-filter-conv2=64 --nb-filter-conv3=128 --nb-filter-conv4=128
---batch-size=64 --dataset=rendered_chairs --epochs=400 --latent_spec_cont=10 --lr=1e-4 --experiment-name=VAE_bs_64_conv_128_128_256_256 --gpu-devices 0 1 --nb-filter-conv1=128 --nb-filter-conv2=128 --nb-filter-conv3=256 --nb-filter-conv4=256
---batch-size=64 --dataset=rendered_chairs --epochs=400 --latent_spec_cont=30 --lr=1e-4 --experiment-name=VAE_bs_64_conv_64_64_128_128_ls_30 --gpu-devices 0 1 --nb-filter-conv1=64 --nb-filter-conv2=64 --nb-filter-conv3=128 --nb-filter-conv4=128
---batch-size=64 --dataset=rendered_chairs --epochs=400 --latent_spec_cont=30 --lr=1e-4 --experiment-name=VAE_bs_64_conv_128_128_256_256_ls_30 --gpu-devices 0 1 --nb-filter-conv1=128 --nb-filter-conv2=128 --nb-filter-conv3=256 --nb-filter-conv4=256

utils/training.py

@@ -63,6 +63,8 @@ class Trainer:
                        'kl_loss': []}
         self.global_iter = 0
         self.mean_epoch_loss = []
+        self.reconstruction_loss = []
+        self.kl_loss = []
         self.save_step = save_step
         self.expe_name = expe_name
         self.ckpt_dir = ckpt_dir
@@ -220,6 +222,7 @@ class Trainer:
 
         # with mse loss:
         recon_loss = F.mse_loss(recon_data, data, size_average=False).div(self.batch_size)
+        self.reconstruction_loss.append(recon_loss)
 
         prediction_loss = 0
         prediction_random_continue_loss = 0
@@ -267,6 +270,7 @@ class Trainer:
 
         # Calculate total kl value to record it
         kl_loss = kl_cont_loss + kl_disc_loss
+        self.kl_loss.append(kl_loss)
 
         # Calculate total loss
         if self.is_beta:
@@ -371,6 +375,8 @@ class Trainer:
         model_states = {'model': self.model.state_dict(), }
         optim_states = {'optim': self.optimizer.state_dict(), }
         states = {'loss': self.mean_epoch_loss,
+                  'recons_loss': self.reconstruction_loss,
+                  'kl_loss': self.kl_loss,
                   'iter': self.global_iter,
                   'model_states': model_states,
                   'optim_states': optim_states}
@@ -385,8 +391,10 @@ class Trainer:
         file_path = os.path.join(self.ckpt_dir, filename)
         if os.path.isfile(file_path):
             checkpoint = torch.load(file_path)
-            if self.load_model_checkpoint:
             self.mean_epoch_loss = checkpoint['loss']
+            if self.load_model_checkpoint:
+                self.reconstruction_loss = checkpoint['recons_loss']
+                self.kl_loss = checkpoint['kl_loss']
             self.global_iter = checkpoint['iter']
             self.model.load_state_dict(checkpoint['model_states']['model'])
             self.optimizer.load_state_dict(checkpoint['optim_states']['optim'])