1
import os
2
from ultralytics import YOLO
3
import torch
4

5
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
6

7
def main():
8
    model_t = YOLO('runs/detect/yolov8x/weights/best.pt')  # the teacher model
9
    model_s = YOLO('runs/detect/yolov8s/weights/best.pt')  # the student model
10
    """
11
    Attributes:
12
        Distillation: the distillation model
13
        loss_type: mgd, cwd
14
        amp: Automatic Mixed Precision
15
    """
16
    model_s.train(data="data.yaml", Distillation=model_t.model, loss_type='mgd', amp=False, imgsz=640, epochs=100,
17
                  batch=20, device=0, workers=0, lr0=0.001)
18

19
if __name__ == '__main__':
20
    main()