Ensure that a checkpoint is saved on the final epoch (#499)

Author: ntw-au

ml3d/tf/pipelines/object_detection.py

@@ -290,7 +290,7 @@ def run_train(self):
 
             self.save_logs(writer, epoch)
 
-            if epoch % cfg.save_ckpt_freq == 0:
+            if epoch % cfg.save_ckpt_freq == 0 or epoch == cfg.max_epoch:
                 self.save_ckpt(epoch)
 
     def get_3d_summary(self,

ml3d/tf/pipelines/semantic_segmentation.py

@@ -306,7 +306,7 @@ def run_train(self):
 
             self.save_logs(writer, epoch)
 
-            if epoch % cfg.save_ckpt_freq == 0:
+            if epoch % cfg.save_ckpt_freq == 0 or epoch == cfg.max_epoch:
                 self.save_ckpt(epoch)
 
     def get_3d_summary(self, results, input_data, epoch, save_gt=True):

ml3d/torch/pipelines/object_detection.py

@@ -346,7 +346,7 @@ def run_train(self):
 
             self.save_logs(writer, epoch)
 
-            if epoch % cfg.save_ckpt_freq == 0:
+            if epoch % cfg.save_ckpt_freq == 0 or epoch == cfg.max_epoch:
                 self.save_ckpt(epoch)
 
     def get_3d_summary(self,

ml3d/torch/pipelines/semantic_segmentation.py

@@ -458,7 +458,7 @@ def run_train(self):
 
             self.save_logs(writer, epoch)
 
-            if epoch % cfg.save_ckpt_freq == 0:
+            if epoch % cfg.save_ckpt_freq == 0 or epoch == cfg.max_epoch:
                 self.save_ckpt(epoch)
 
     def get_batcher(self, device, split='training'):