Model enhancement (#75)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/tests/test_model.py

@@ -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()