Algorithm-Repo/SourceCode/MachineLearningUtils/VINNOFramework/YOLO/main.py

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':

    # model = YOLO(model='yolov8n-seg.yaml')
    # # model.load('yolov8n.pt') # loading pretrain weights
    # model.train(data='vinno_dataset/Liver-Focal-Seg.yaml',
    #             cache=False,
    #             imgsz=320,
    #             epochs=2,
    #             batch=16,
    #             close_mosaic=0,
    #             workers=1,
    #             device='0',
    #             optimizer='SGD', # using SGD
    #             # patience=0, # close earlystop
    #             # resume='', # last.pt path
    #             # amp=False, # close amp
    #             # fraction=0.2,
    #             project='runs/train',
    #             name='exp',
    #             )

    from ultralytics import RTDETR
    model = RTDETR("rt-detr/rtdetr-l.yaml")
    # Display model information (optional)
    model.info()
    # Train the model on the COCO8 example dataset for 100 epochs
    results = model.train(data='vinno_dataset/Thyroid-TIRADS-Obj.yaml',
                cache=False,
                imgsz=320,
                epochs=2,
                batch=4,
                close_mosaic=0,
                workers=1,
                device='0',
                optimizer='SGD', # using SGD
                # patience=0, # close earlystop
                # resume='', # last.pt path
                # amp=False, # close amp
                # fraction=0.2,
                project='runs/train',
                name='exp',)