naseemap47 / YOLO-NAS

Train and Inference your custom YOLO-NAS model by Single Command Line
Apache License 2.0
98 stars 13 forks source link

Resume #45

Closed naseemap47 closed 1 year ago

naseemap47 commented 1 year ago

To solve #44 issue. added resume option for to continue training from last ckpt.

🤖 Train

You can train your YOLO-NAS model with Single Command Line

Args `-i`, `--data`: path to data.yaml `-n`, `--name`: Checkpoint dir name `-b`, `--batch`: Training batch size `-e`, `--epoch`: number of training epochs. `-s`, `--size`: Input image size `-j`, `--worker`: Training number of workers `-m`, `--model`: Model type (Choices: `yolo_nas_s`, `yolo_nas_m`, `yolo_nas_l`) `-w`, `--weight`: path to pre-trained model weight (default: `coco` weight) `--gpus`: Train on multiple gpus `--cpu`: Train on CPU `--resume`: To resume model training **Other Training Parameters:** `--warmup_mode`: Warmup Mode, eg: Linear Epoch Step `--warmup_initial_lr`: Warmup Initial LR `--lr_warmup_epochs`: LR Warmup Epochs `--initial_lr`: Inital LR `--lr_mode`: LR Mode, eg: cosine `--cosine_final_lr_ratio`: Cosine Final LR Ratio `--optimizer`: Optimizer, eg: Adam `--weight_decay`: Weight Decay

Example:

python3 train.py --data /dir/dataset/data.yaml --batch 6 --epoch 100 --model yolo_nas_m --size 640