Merge pull request #45 from naseemap47/resume

Resume

README.md

@@ -116,6 +116,7 @@ You can train your **YOLO-NAS** model with **Single Command Line**
   `-w`, `--weight`: path to pre-trained model weight (default: `coco` weight) <br>
   `--gpus`: Train on multiple gpus <br>
   `--cpu`: Train on CPU <br>
+  `--resume`: To resume model training <br>
 
   **Other Training Parameters:**<br>
   `--warmup_mode`: Warmup Mode, eg: Linear Epoch Step <br>

train.py

@@ -29,6 +29,8 @@
                 help="Model type (eg: yolo_nas_s)")
     ap.add_argument("-w", "--weight", type=str, default='coco',
                     help="path to pre-trained model weight")
+    ap.add_argument("-r", "--resume", action='store_true',
+                    help="to resume model training")
     ap.add_argument("-s", "--size", type=int, default=640,
                     help="input image size")
     ap.add_argument("--gpus", action='store_true',
@@ -128,11 +130,19 @@
             }
         )
 
-    model = models.get(
-        args['model'],
-        num_classes=len(yaml_params['names']), 
-        pretrained_weights=args["weight"]
-    )
+    # To Resume Training
+    if args['resume']:
+        model = models.get(
+            args['model'],
+            num_classes=len(yaml_params['names']),
+            checkpoint_path=args["weight"]
+        )
+    else:
+        model = models.get(
+            args['model'],
+            num_classes=len(yaml_params['names']), 
+            pretrained_weights=args["weight"]
+        )
 
     train_params = {
         # ENABLING SILENT MODE