diff --git a/tests/test_python.py b/tests/test_python.py index 7f0dc7c3..67223c05 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -640,26 +640,31 @@ def test_yolo_world(): model.set_classes(["tree", "window"]) model(ASSETS / "bus.jpg", conf=0.01) - # Training from yaml - model = YOLO("yolov8s-worldv2.yaml") # no YOLOv8n-world model yet - model.train(data="coco8.yaml", epochs=2, imgsz=32, cache="disk", batch=-1, close_mosaic=1, name="yolo-world") - model = YOLO("yolov8s-worldv2.pt") # no YOLOv8n-world model yet - # val - model.val(data="coco8.yaml", imgsz=32, save_txt=True, save_json=True) - # Training from pretrain - model.train(data="coco8.yaml", epochs=2, imgsz=32, cache="disk", batch=-1, close_mosaic=1, name="yolo-world") + # Training from pretrain, evaluation process is included at the final stage of training. + # Use dota8.yaml which has less categories to reduce the inference time of CLIP model + model.train( + data="dota8.yaml", + epochs=1, + imgsz=32, + cache="disk", + batch=4, + close_mosaic=1, + name="yolo-world", + save_txt=True, + save_json=True, + ) # test WorWorldTrainerFromScratch from ultralytics.models.yolo.world.train_world import WorldTrainerFromScratch model = YOLO("yolov8s-worldv2.yaml") # no YOLOv8n-world model yet model.train( - data={"train": {"yolo_data": ["coco8.yaml"]}, "val": {"yolo_data": ["coco8.yaml"]}}, - epochs=2, + data={"train": {"yolo_data": ["dota8.yaml"]}, "val": {"yolo_data": ["dota8.yaml"]}}, + epochs=1, imgsz=32, cache="disk", - batch=-1, + batch=4, close_mosaic=1, name="yolo-world", trainer=WorldTrainerFromScratch,