From 1f654186c512fd9b53c19ea4d93bc3f01583e5fe Mon Sep 17 00:00:00 2001
From: NASEEM A P <naseemalassampattil@gmail.com>
Date: Fri, 22 Sep 2023 10:34:26 +0530
Subject: [PATCH 1/2] added resume option for to continue training from last
 ckpt

---
 train.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/train.py b/train.py
index 5565b02..896d093 100644
--- a/train.py
+++ b/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

From 357921c226a44d890eb6e0135d51a5ea43d3821e Mon Sep 17 00:00:00 2001
From: NASEEM A P <naseemalassampattil@gmail.com>
Date: Fri, 22 Sep 2023 10:36:21 +0530
Subject: [PATCH 2/2] added train args - resume

---
 README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/README.md b/README.md
index d8ce695..10c4638 100644
--- a/README.md
+++ b/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>