File size: 1,556 Bytes
9629250 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
from ultralytics import YOLO
def train_model_and_store(model_name, epochs, img_size, batch_size, device, optimizer, learning_rate):
# Load the pretrained model
model = YOLO(f'./base_models/{model_name}')
model.train(
data='./BracU_Ds/data.yaml',
epochs=epochs,
imgsz=img_size,
batch=batch_size,
device=device,
project='Activity_Detection_New_BracU',
name=model_name,
optimizer=optimizer,
lr0=learning_rate,
patience=10,
plots=True,
seed=42,
# pretrained=True,
# degrees=15,
# translate=0.2,
# scale=0.8,
# shear=10.0,
# perspective=0.001,
# fliplr=0.5,
# mosaic=1.0,
# mixup=0.2,
# hsv_h=0.015,
# hsv_s=0.7,
# hsv_v=0.4
)
# Load the best model after training
# best_model_path = f'./Activity_Detection_New_BracU/{model_name}/weights/best.pt'
# model = YOLO(best_model_path)
# Evaluate the model
# model.val(project='Activity_Detection_New_Eval', name=model_name)
if __name__ == "__main__":
train_model_and_store(
model_name="yolov8n.pt",
epochs=50,
img_size=640,
batch_size=64,
device=0,
optimizer='Adam',
learning_rate=0.005
)
# We should fine tune the following model
# yolo11n.pt
# yolov10n.pt
# yolov8n.pt
# yolov8m.pt
# yolo11n-cls.pt
# yolov8n-cls.pt
# nohup python train.py > logs/training_log_yolov8n_BracU_Ds.txt 2>&1 &
# 2980863
|