yolo training on custom dataset

requirements.txt

@@ -1,2 +1,3 @@
 dvc==2.51.0
+ultralytics==8.0.58
 

src/params.yaml

@@ -0,0 +1,9 @@
+model_type: yolov8n.pt
+pretrained: True
+seed: 0
+imgsz: 640 
+batch: 8 
+epochs: 1
+optimizer: SGD # other choices=['SGD', 'Adam', 'AdamW', 'RMSProp']
+lr0: 0.01 # learning rate 
+name: 'yolov8n_exp_v0' # experiment name

src/train.py

@@ -0,0 +1,42 @@
+from ultralytics import YOLO
+import yaml
+from pathlib import Path
+import os
+
+
+ROOT_DIR = Path(__file__).resolve().parents[1] # root directory absolute path
+DATA_DIR = os.path.join(ROOT_DIR, "data/raw/wildfire-raw-yolov8")
+DATA_YAML = os.path.join(DATA_DIR, "data.yaml")
+print(DATA_YAML)
+
+
+if __name__ == '__main__':
+
+ with open(r"src/params.yaml") as f:
+ params = yaml.safe_load(f)
+
+ # load a pre-trained model 
+ pre_trained_model = YOLO(params['model_type'])
+
+ # train 
+ model = pre_trained_model.train(
+ data=DATA_YAML,
+ imgsz=params['imgsz'],
+ batch=params['batch'],
+ epochs=params['epochs'],
+ optimizer=params['optimizer'],
+ lr0=params['lr0'],
+ seed=params['seed'],
+ pretrained=params['pretrained'],
+ name=params['name']
+ )
+
+ # evaluate model 
+ model.val()
+
+
+
+
+
+
+