Model interface enhancement (#106)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia
2022-12-28 18:05:01 +05:30
committed by GitHub
parent 38d6df55cb
commit 384f0ef1c6
6 changed files with 39 additions and 22 deletions

View File

@ -1,11 +1,11 @@
import torch
from ultralytics.yolo import YOLO
from ultralytics import YOLO
def test_model_forward():
model = YOLO()
model.new("yolov8n-seg.yaml")
model.new("yolov8n.yaml")
img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512)
model.forward(img)
model(img)
@ -15,7 +15,7 @@ def test_model_info():
model = YOLO()
model.new("yolov8n.yaml")
model.info()
model.load("balloon-detect.pt")
model.load("best.pt")
model.info(verbose=True)
@ -23,35 +23,35 @@ def test_model_fuse():
model = YOLO()
model.new("yolov8n.yaml")
model.fuse()
model.load("balloon-detect.pt")
model.load("best.pt")
model.fuse()
def test_visualize_preds():
model = YOLO()
model.load("balloon-segment.pt")
model.load("best.pt")
model.predict(source="ultralytics/assets")
def test_val():
model = YOLO()
model.load("balloon-segment.pt")
model.val(data="coco128-seg.yaml", imgsz=32)
model.load("best.pt")
model.val(data="coco128.yaml", imgsz=32)
def test_model_resume():
model = YOLO()
model.new("yolov8n-seg.yaml")
model.train(epochs=1, imgsz=32, data="coco128-seg.yaml")
model.new("yolov8n.yaml")
model.train(epochs=1, imgsz=32, data="coco128.yaml")
try:
model.resume(task="segment")
model.resume(task="detect")
except AssertionError:
print("Successfully caught resume assert!")
def test_model_train_pretrained():
model = YOLO()
model.load("balloon-detect.pt")
model.load("best.pt")
model.train(data="coco128.yaml", epochs=1, imgsz=32)
model.new("yolov8n.yaml")
model.train(data="coco128.yaml", epochs=1, imgsz=32)