Model enhancement (#75)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
@ -1,13 +1,62 @@
|
||||
import torch
|
||||
|
||||
from ultralytics.yolo import YOLO
|
||||
|
||||
|
||||
def test_model():
|
||||
def test_model_forward():
|
||||
model = YOLO()
|
||||
model.new("assets/dummy_model.yaml")
|
||||
model.model = "squeezenet1_0" # temp solution before get_model is implemented
|
||||
# model.load("yolov5n.pt")
|
||||
model.train(data="imagenette160", epochs=1, lr0=0.01)
|
||||
model.new("yolov5n-seg.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("yolov5n.yaml")
|
||||
model.info()
|
||||
model.load("balloon-detect.pt")
|
||||
model.info(verbose=True)
|
||||
|
||||
|
||||
def test_model_fuse():
|
||||
model = YOLO()
|
||||
model.new("yolov5n.yaml")
|
||||
model.fuse()
|
||||
model.load("balloon-detect.pt")
|
||||
model.fuse()
|
||||
|
||||
|
||||
def test_visualize_preds():
|
||||
model = YOLO()
|
||||
model.load("balloon-segment.pt")
|
||||
model.predict(source="ultralytics/assets")
|
||||
|
||||
|
||||
def test_val():
|
||||
model = YOLO()
|
||||
model.load("balloon-segment.pt")
|
||||
model.val(data="coco128-seg.yaml", img_size=32)
|
||||
|
||||
|
||||
def test_model_resume():
|
||||
model = YOLO()
|
||||
model.new("yolov5n-seg.yaml")
|
||||
model.train(epochs=1, img_size=32, data="coco128-seg.yaml")
|
||||
try:
|
||||
model.resume(task="segment")
|
||||
except AssertionError:
|
||||
print("Successfully caught resume assert!")
|
||||
|
||||
|
||||
def test():
|
||||
test_model_forward()
|
||||
test_model_info()
|
||||
test_model_fuse()
|
||||
test_visualize_preds()
|
||||
test_val()
|
||||
test_model_resume()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_model()
|
||||
test()
|
||||
|
Reference in New Issue
Block a user