optimize saveckpt callback

Author: Mark-ZhouWX

research/segment-anything/README.md

@@ -79,17 +79,17 @@ See `python use_sam_with_amg.py --help` to explore more custom settings.
 ## Finetune
 
 Finetune is a popular method that adapts large pretrained model to specific downstream tasks. Currently, finetune with box-prompt are supported. The bounding boxes are used as prompt input to predict mask.
-Beside fine-tuning our code on COCO2017 dataset which contains common seen objects and lies in the similar distribution of the original [training dataset](https://segment-anything.com/dataset/index.html)) of SAM, We have done further experiments on a medical imaging segmentation dataset [FLARE22](https://flare22.grand-challenge.org/Dataset/). Result shows that the finetune method in this repository is effective.
+Beside fine-tuning our code on COCO2017 dataset which contains common seen objects and lies in the similar distribution of the original [training dataset](https://segment-anything.com/dataset/index.html) of SAM, We have done further experiments on a medical imaging segmentation dataset [FLARE22](https://flare22.grand-challenge.org/Dataset/). Result shows that the finetune method in this repository is effective.
 
 The bellowing shows the mask quality before and after finetune.
 
 
-| pretrained_model | dataset  |    epochs    | mIOU |
-| :--------------: | -------- | :-----------: | ---- |
-|    sam-vit-b    | COCO2017 | 0 (zero-shot) | 77.4 |
-|    sam-vit-b    | COCO2017 |      20      | 83.6 |
-|    sam-vit-b    | FLARE22  | 0 (zero-shot) | 79.5 |
-|    sam-vit-b    | FLARE22  |      10      | 88.1 |
+| pretrained_model | dataset  |    epochs     | mIOU |
+|:----------------:| -------- |:-------------:|------|
+|    sam-vit-b     | COCO2017 | 0 (zero-shot) | 77.4 |
+|    sam-vit-b     | COCO2017 |      20       | 83.5 |
+|    sam-vit-b     | FLARE22  | 0 (zero-shot) | 79.5 |
+|    sam-vit-b     | FLARE22  |      10       | 88.1 |
 
 To finetune COCO dataset, please run:
 

research/segment-anything/configs/coco_box_finetune.yaml

@@ -100,6 +100,7 @@ eval_metric: &eval_metric
 callback:
   - type: segment_anything.utils.callbacks.TrainStatusLog
     loss_item: ['focal_loss', 'dice_loss', 'mse_loss']  # for log
+    interval: 100
   - type: segment_anything.utils.callbacks.SaveCkpt
     work_root: *work_root
     interval: 1  # in epoch

research/segment-anything/configs/flare_box_finetune.yaml

@@ -95,7 +95,7 @@ eval_metric: &eval_metric
 callback:
   - type: segment_anything.utils.callbacks.TrainStatusLog
     loss_item: ['focal_loss', 'dice_loss', 'mse_loss']  # for log
-    interval: 20
+    interval: 100
   - type: segment_anything.utils.callbacks.SaveCkpt
     work_root: *work_root
     interval: 1  # in epoch

research/segment-anything/segment_anything/utils/callbacks.py

@@ -83,7 +83,7 @@ class TrainStatusLog(Callback):
     """
     Callback to record the status of training, mainly including loss and time performance information.
     """
-    def __init__(self, interval=10, loss_item=()):
+    def __init__(self, interval=100, loss_item=()):
         self.log_interval = interval
         self.loss_item = loss_item
         self.step_start_time = 0.0
@@ -166,7 +166,8 @@ def __init__(self, interval=1, work_root='./work_root', save_dir='', main_device
     def on_train_epoch_end(self, run_context: RunContext):
         cb_params = run_context.original_args()
         cur_epoch = cb_params.cur_epoch_num
-        if self.main_device and cur_epoch % self.interval == 0:
+        total_epoch_num = cb_params.epoch_num
+        if self.main_device and (cur_epoch % self.interval == 0 or cur_epoch == total_epoch_num):
             save_path = os.path.join(self.full_save_dir, f'sam_{cur_epoch:03d}.ckpt')
             logger.info(f'saving ckpt of epoch {cur_epoch} at {save_path}, interval is {self.interval}')
             # model without loss function, cb_params.network is train_one_step_cell