Make YOLO a module (#111)
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:
@ -3,45 +3,49 @@ import torch
|
||||
from ultralytics import YOLO
|
||||
|
||||
|
||||
def test_model_init():
|
||||
model = YOLO.new("yolov8n.yaml")
|
||||
model.info()
|
||||
try:
|
||||
YOLO()
|
||||
except Exception:
|
||||
print("Successfully caught constructor assert!")
|
||||
raise Exception("constructor error didn't occur")
|
||||
|
||||
|
||||
def test_model_forward():
|
||||
model = YOLO()
|
||||
model.new("yolov8n.yaml")
|
||||
model = YOLO.new("yolov8n.yaml")
|
||||
img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512)
|
||||
model.forward(img)
|
||||
model(img)
|
||||
|
||||
|
||||
def test_model_info():
|
||||
model = YOLO()
|
||||
model.new("yolov8n.yaml")
|
||||
model = YOLO.new("yolov8n.yaml")
|
||||
model.info()
|
||||
model.load("best.pt")
|
||||
model = model.load("best.pt")
|
||||
model.info(verbose=True)
|
||||
|
||||
|
||||
def test_model_fuse():
|
||||
model = YOLO()
|
||||
model.new("yolov8n.yaml")
|
||||
model = YOLO.new("yolov8n.yaml")
|
||||
model.fuse()
|
||||
model.load("best.pt")
|
||||
model.fuse()
|
||||
|
||||
|
||||
def test_visualize_preds():
|
||||
model = YOLO()
|
||||
model.load("best.pt")
|
||||
model = YOLO.load("best.pt")
|
||||
model.predict(source="ultralytics/assets")
|
||||
|
||||
|
||||
def test_val():
|
||||
model = YOLO()
|
||||
model.load("best.pt")
|
||||
model = YOLO.load("best.pt")
|
||||
model.val(data="coco128.yaml", imgsz=32)
|
||||
|
||||
|
||||
def test_model_resume():
|
||||
model = YOLO()
|
||||
model.new("yolov8n.yaml")
|
||||
model = YOLO.new("yolov8n.yaml")
|
||||
model.train(epochs=1, imgsz=32, data="coco128.yaml")
|
||||
try:
|
||||
model.resume(task="detect")
|
||||
@ -50,10 +54,9 @@ def test_model_resume():
|
||||
|
||||
|
||||
def test_model_train_pretrained():
|
||||
model = YOLO()
|
||||
model.load("best.pt")
|
||||
model = YOLO.load("best.pt")
|
||||
model.train(data="coco128.yaml", epochs=1, imgsz=32)
|
||||
model.new("yolov8n.yaml")
|
||||
model = model.new("yolov8n.yaml")
|
||||
model.train(data="coco128.yaml", epochs=1, imgsz=32)
|
||||
img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512)
|
||||
model(img)
|
||||
|
Reference in New Issue
Block a user