save all

scripts/cvae_pyramid_e2e.py

@@ -547,7 +547,7 @@ def train():
         valid_accumulator.update(valid_loss)
 
         if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
-            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step, 'balancing_coeff': g_balancing_coeff } })
             if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
                 train_figs = plot_everything(batch, 'train')
                 valid_figs = plot_everything(data_valid, 'valid')
@@ -561,10 +561,9 @@ def train():
             valid_accumulator.reset()
             print(step)
 
-        # if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
-        #     g_balancing_coeff += 1 / 50000 
+        g_balancing_coeff += 1 / 100000 
 
-        if step % 5000 == 0 and not args['debug']:
+        if step % 10000 == 0 and not args['debug']:
             ckpt_saver(model, step, optimizer, None)
 
         step += 1
@@ -593,7 +592,7 @@ if __name__ == '__main__':
     # optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=20)
     scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
     resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])