update the model weights

README.md

@@ -6,19 +6,25 @@ library_name: monai
 license: apache-2.0
 ---
 # Model Overview
-A pre-trained model for the endoscopic tool segmentation task and is trained using a flexible unet structure with an efficient-b2 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
+A pre-trained model for the endoscopic tool segmentation task, trained using a flexible unet structure with an efficientnet-b2 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
 
-The [PyTorch model](https://drive.google.com/file/d/19yS3t2oLBiB7wT-qeQ82da95VJs_vzRK/view?usp=share_link) and [torchscript model](https://drive.google.com/file/d/1cDZ3Jr7mhpzdzaFyz8yHNowH8k0T1VZz/view?usp=share_link) are shared in google drive. Details can be found in large_files.yml file. Modify the "bundle_root" parameter specified in configs/train.json and configs/inference.json to reflect where models are downloaded. Expected directory path to place downloaded models is "models/" under "bundle_root".
+The [PyTorch model](https://drive.google.com/file/d/1I7UtWDKDEcezMqYiA-i_hsRTCrvWwJ61/view?usp=sharing) and [torchscript model](https://drive.google.com/file/d/1e_wYd1HjJQ0dz_HKdbthRcMOyUL02aLG/view?usp=sharing) are shared in google drive. Details can be found in `large_files.yml` file. Modify the "bundle_root" parameter specified in configs/train.json and configs/inference.json to reflect where models are downloaded. Expected directory path to place downloaded models is "models/" under "bundle_root".
 
 ![image](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_workflow.png)
 
+## Pre-trained weights
+A pre-trained encoder weights would benefit the model training. In this bundle, the encoder is trained with pre-trained weights from some internal data. For convenience, we provide two options in `configs/train.json` to enable users to load pre-trained weights:
+
+1. Via setting the `use_imagenet_pretrain` parameter in the config file to `True`, [ImageNet](https://ieeexplore.ieee.org/document/5206848) pre-trained weights from the [EfficientNet-PyTorch repo](https://github.com/lukemelas/EfficientNet-PyTorch) can be loaded. Please note that these weights are for non-commercial use. Each user is responsible for checking the content of the models/datasets and the applicable licenses and determining if suitable for the intended use.
+2. Via updating the `load_path`  parameter of the `CheckpointLoader` in the config file, weights from a local path can be loaded.
+
 ## Data
 Datasets used in this work were provided by [Activ Surgical](https://www.activsurgical.com/).
 
 Since datasets are private, existing public datasets like [EndoVis 2017](https://endovissub2017-roboticinstrumentsegmentation.grand-challenge.org/Data/) can be used to train a similar model.
 
 ### Preprocessing
-When using EndoVis or any other dataset, it should be divided into "train", "valid" and "test" folders. Samples in each folder would better be images and converted to jpg format. Otherwise, "images", "labels", "val_images" and "val_labels" parameters in "configs/train.json" and "datalist" in "configs/inference.json" should be modified to fit given dataset. After that, "dataset_dir" parameter in "configs/train.json" and "configs/inference.json" should be changed to root folder which contains previous "train", "valid" and "test" folders.
+When using EndoVis or any other dataset, it should be divided into "train", "valid" and "test" folders. Samples in each folder would better be images and converted to jpg format. Otherwise, "images", "labels", "val_images" and "val_labels" parameters in `configs/train.json` and "datalist" in `configs/inference.json` should be modified to fit given dataset. After that, "dataset_dir" parameter in `configs/train.json` and `configs/inference.json` should be changed to root folder which contains previous "train", "valid" and "test" folders.
 
 Please notice that loading data operation in this bundle is adaptive. If images and labels are not in the same format, it may lead to a mismatching problem. For example, if images are in jpg format and labels are in npy format, PIL and Numpy readers will be used separately to load images and labels. Since these two readers have their own way to parse file's shape, loaded labels will be transpose of the correct ones and incur a missmatching problem.
 
@@ -38,13 +44,13 @@ Two channels:
 - Label 0: everything else
 
 ## Performance
-IoU was used for evaluating the performance of the model. This model achieves a mean IoU score of 0.87.
+IoU was used for evaluating the performance of the model. This model achieves a mean IoU score of 0.86.
 
 #### Training Loss
-![A graph showing the training loss over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_train_loss.png)
+![A graph showing the training loss over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_train_loss_v3.png)
 
 #### Validation IoU
-![A graph showing the validation mean IoU over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_val_iou.png)
+![A graph showing the validation mean IoU over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_val_iou_v3.png)
 
 #### TensorRT speedup
 The `endoscopic_tool_segmentation` bundle supports the TensorRT acceleration. The table below shows the speedup ratios benchmarked on an A100 80G GPU.

configs/evaluate.json

@@ -47,7 +47,8 @@
             "load_path": "$@ckpt_dir + '/model.pt'",
             "load_dict": {
                 "model": "@network"
-            }
+            },
+            "map_location": "@device"
         },
         {
             "_target_": "StatsHandler",
@@ -68,9 +69,6 @@
             "summary_ops": "*"
         }
     ],
-    "initialize": [
-        "$setattr(torch.backends.cudnn, 'benchmark', True)"
-    ],
     "run": [
         "$@validate#evaluator.run()"
     ]

configs/inference.json

@@ -114,7 +114,8 @@
             "load_path": "$@bundle_root + '/models/model.pt'",
             "load_dict": {
                 "model": "@network"
-            }
+            },
+            "map_location": "@device"
         },
         {
             "_target_": "StatsHandler",
@@ -131,7 +132,7 @@
         "val_handlers": "@handlers"
     },
     "initialize": [
-        "$setattr(torch.backends.cudnn, 'benchmark', True)"
+        "$monai.utils.set_determinism(seed=123)"
     ],
     "run": [
         "$@evaluator.run()"

configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json",
-    "version": "0.4.8",
+    "version": "0.4.9",
     "changelog": {
+        "0.4.9": "update the model weights",
         "0.4.8": "update the TensorRT part in the README file",
         "0.4.7": "fix mgpu finalize issue",
         "0.4.6": "enable deterministic training",
@@ -37,7 +38,7 @@
     "label_classes": "single channel data, 1/255 is tool, 0 is background",
     "pred_classes": "2 channels OneHot data, channel 1 is tool, channel 0 is background",
     "eval_metrics": {
-        "mean_iou": 0.87
+        "mean_iou": 0.86
     },
     "references": [
         "Tan, M. and Le, Q. V. Efficientnet: Rethinking model scaling for convolutional neural networks. ICML, 2019a. https://arxiv.org/pdf/1905.11946.pdf",

configs/train.json

@@ -6,6 +6,7 @@
         "$import numpy as np"
     ],
     "bundle_root": ".",
+    "use_imagenet_pretrain": false,
     "ckpt_dir": "$@bundle_root + '/models'",
     "output_dir": "$@bundle_root + '/eval'",
     "dataset_dir": "/workspace/data/endoscopic_tool_dataset",
@@ -21,14 +22,14 @@
         "out_channels": 2,
         "backbone": "efficientnet-b2",
         "spatial_dims": 2,
-        "dropout": 0.5,
-        "pretrained": true,
+        "dropout": 0.6,
+        "pretrained": "@use_imagenet_pretrain",
         "is_pad": false,
         "pre_conv": null
     },
     "network": "$@network_def.to(@device)",
     "loss": {
-        "_target_": "DiceLoss",
+        "_target_": "DiceFocalLoss",
         "include_background": false,
         "to_onehot_y": true,
         "softmax": true,
@@ -122,7 +123,7 @@
             "_target_": "CacheDataset",
             "data": "$[{'image': i, 'label': l} for i, l in zip(@images, @labels)]",
             "transform": "@train#preprocessing",
-            "cache_rate": 0.1,
+            "cache_rate": 0.5,
             "num_workers": 4
         },
         "dataloader": {
@@ -136,6 +137,16 @@
             "_target_": "SimpleInferer"
         },
         "handlers": [
+            {
+                "_target_": "CheckpointLoader",
+                "_disabled_": "@use_imagenet_pretrain",
+                "load_path": "/path/to/local/weight/model.pt",
+                "load_dict": {
+                    "model": "@network"
+                },
+                "strict": false,
+                "map_location": "@device"
+            },
             {
                 "_target_": "ValidationHandler",
                 "validator": "@validate#evaluator",
@@ -210,7 +221,7 @@
             "_target_": "CacheDataset",
             "data": "$[{'image': i, 'label': l} for i, l in zip(@val_images, @val_labels)]",
             "transform": "@validate#preprocessing",
-            "cache_rate": 0.1
+            "cache_rate": 1.0
         },
         "dataloader": {
             "_target_": "DataLoader",

docs/README.md

@@ -1,17 +1,23 @@
 # Model Overview
-A pre-trained model for the endoscopic tool segmentation task and is trained using a flexible unet structure with an efficient-b2 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
+A pre-trained model for the endoscopic tool segmentation task, trained using a flexible unet structure with an efficientnet-b2 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
 
-The [PyTorch model](https://drive.google.com/file/d/19yS3t2oLBiB7wT-qeQ82da95VJs_vzRK/view?usp=share_link) and [torchscript model](https://drive.google.com/file/d/1cDZ3Jr7mhpzdzaFyz8yHNowH8k0T1VZz/view?usp=share_link) are shared in google drive. Details can be found in large_files.yml file. Modify the "bundle_root" parameter specified in configs/train.json and configs/inference.json to reflect where models are downloaded. Expected directory path to place downloaded models is "models/" under "bundle_root".
+The [PyTorch model](https://drive.google.com/file/d/1I7UtWDKDEcezMqYiA-i_hsRTCrvWwJ61/view?usp=sharing) and [torchscript model](https://drive.google.com/file/d/1e_wYd1HjJQ0dz_HKdbthRcMOyUL02aLG/view?usp=sharing) are shared in google drive. Details can be found in `large_files.yml` file. Modify the "bundle_root" parameter specified in configs/train.json and configs/inference.json to reflect where models are downloaded. Expected directory path to place downloaded models is "models/" under "bundle_root".
 
 ![image](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_workflow.png)
 
+## Pre-trained weights
+A pre-trained encoder weights would benefit the model training. In this bundle, the encoder is trained with pre-trained weights from some internal data. For convenience, we provide two options in `configs/train.json` to enable users to load pre-trained weights:
+
+1. Via setting the `use_imagenet_pretrain` parameter in the config file to `True`, [ImageNet](https://ieeexplore.ieee.org/document/5206848) pre-trained weights from the [EfficientNet-PyTorch repo](https://github.com/lukemelas/EfficientNet-PyTorch) can be loaded. Please note that these weights are for non-commercial use. Each user is responsible for checking the content of the models/datasets and the applicable licenses and determining if suitable for the intended use.
+2. Via updating the `load_path`  parameter of the `CheckpointLoader` in the config file, weights from a local path can be loaded.
+
 ## Data
 Datasets used in this work were provided by [Activ Surgical](https://www.activsurgical.com/).
 
 Since datasets are private, existing public datasets like [EndoVis 2017](https://endovissub2017-roboticinstrumentsegmentation.grand-challenge.org/Data/) can be used to train a similar model.
 
 ### Preprocessing
-When using EndoVis or any other dataset, it should be divided into "train", "valid" and "test" folders. Samples in each folder would better be images and converted to jpg format. Otherwise, "images", "labels", "val_images" and "val_labels" parameters in "configs/train.json" and "datalist" in "configs/inference.json" should be modified to fit given dataset. After that, "dataset_dir" parameter in "configs/train.json" and "configs/inference.json" should be changed to root folder which contains previous "train", "valid" and "test" folders.
+When using EndoVis or any other dataset, it should be divided into "train", "valid" and "test" folders. Samples in each folder would better be images and converted to jpg format. Otherwise, "images", "labels", "val_images" and "val_labels" parameters in `configs/train.json` and "datalist" in `configs/inference.json` should be modified to fit given dataset. After that, "dataset_dir" parameter in `configs/train.json` and `configs/inference.json` should be changed to root folder which contains previous "train", "valid" and "test" folders.
 
 Please notice that loading data operation in this bundle is adaptive. If images and labels are not in the same format, it may lead to a mismatching problem. For example, if images are in jpg format and labels are in npy format, PIL and Numpy readers will be used separately to load images and labels. Since these two readers have their own way to parse file's shape, loaded labels will be transpose of the correct ones and incur a missmatching problem.
 
@@ -31,13 +37,13 @@ Two channels:
 - Label 0: everything else
 
 ## Performance
-IoU was used for evaluating the performance of the model. This model achieves a mean IoU score of 0.87.
+IoU was used for evaluating the performance of the model. This model achieves a mean IoU score of 0.86.
 
 #### Training Loss
-![A graph showing the training loss over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_train_loss.png)
+![A graph showing the training loss over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_train_loss_v3.png)
 
 #### Validation IoU
-![A graph showing the validation mean IoU over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_val_iou.png)
+![A graph showing the validation mean IoU over 100 epochs.](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_val_iou_v3.png)
 
 #### TensorRT speedup
 The `endoscopic_tool_segmentation` bundle supports the TensorRT acceleration. The table below shows the speedup ratios benchmarked on an A100 80G GPU.

models/model.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:844e9e97c6c9e7ebab1dab660c42ceb923c085ab895cf74f989a3cd9c5a0b028
-size 46262677
+oid sha256:1a24d5f1d30a44474e8bb87abb62f289add2c40a291eac67138242679e928729
+size 46290573

models/model.ts

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9fa9f747ea2cc8ddffd4839af0c8d8f1b62c63c1013a8d80af8baf412bb3e5f9
-size 46493609
+oid sha256:149ac9473dee60737957505a59a24c6d816092c6aa2ff89e13c8bfc8e9831059
+size 46500237