Add multiple checkpoints for different model and set default config value

1 files changed, 9 insertions, 5 deletions

diff --git a/config.py b/config.py
index 8a8d93a..7517097 100644
--- a/config.py
+++ b/config.py
@@ -13,7 +13,7 @@ config: Configuration = {
     'dataset': {
         # Name of dataset (CASIA-B or FVG)
         'name': 'CASIA-B',
-        # Path to dataset root
+        # Path to dataset root (required)
         'root_dir': 'data/CASIA-B-MRCNN/SEG',
         # The number of subjects for training
         'train_size': 74,
@@ -88,9 +88,13 @@ config: Configuration = {
     'model': {
         # Model name, used for naming checkpoint
         'name': 'RGB-GaitPart',
-        # Restoration iteration from checkpoint
-        'restore_iter': 0,
-        # Total iteration for training
-        'total_iter': 80000,
+        # Restoration iteration from checkpoint (single model)
+        # 'restore_iter': 0,
+        # Total iteration for training (single model)
+        # 'total_iter': 80000,
+        # Restoration iteration (multiple models, e.g. nm, bg and cl)
+        'restore_iters': (0, 0, 0),
+        # Total iteration for training (multiple models)
+        'total_iter': (80_000, 80_000, 80_000),
     },
 }