You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
100 lines
2.7 KiB
100 lines
2.7 KiB
import torch
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
def test_model_init():
|
|
model = YOLO("yolov8n.yaml")
|
|
model.info()
|
|
|
|
|
|
def test_model_forward():
|
|
model = YOLO("yolov8n.yaml")
|
|
img = torch.rand(1, 3, 320, 320)
|
|
model.forward(img)
|
|
model(img)
|
|
|
|
|
|
def test_model_info():
|
|
model = YOLO("yolov8n.yaml")
|
|
model.info()
|
|
model = YOLO("yolov8n.pt")
|
|
model.info(verbose=True)
|
|
|
|
|
|
def test_model_fuse():
|
|
model = YOLO("yolov8n.yaml")
|
|
model.fuse()
|
|
model = YOLO("yolov8n.pt")
|
|
model.fuse()
|
|
|
|
|
|
def test_visualize_preds():
|
|
model = YOLO("yolov8n.pt")
|
|
model.predict(source="ultralytics/assets")
|
|
|
|
|
|
def test_val():
|
|
model = YOLO("yolov8n.pt")
|
|
model.val(data="coco128.yaml", imgsz=32)
|
|
|
|
|
|
def test_model_resume():
|
|
model = YOLO("yolov8n.yaml")
|
|
model.train(epochs=1, imgsz=32, data="coco128.yaml")
|
|
try:
|
|
model.resume(task="detect")
|
|
except AssertionError:
|
|
print("Successfully caught resume assert!")
|
|
|
|
|
|
def test_model_train_pretrained():
|
|
model = YOLO("yolov8n.pt")
|
|
model.train(data="coco128.yaml", epochs=1, imgsz=32)
|
|
model = YOLO("yolov8n.yaml")
|
|
model.train(data="coco128.yaml", epochs=1, imgsz=32)
|
|
img = torch.rand(1, 3, 320, 320)
|
|
model(img)
|
|
|
|
|
|
def test_exports():
|
|
"""
|
|
Format Argument Suffix CPU GPU
|
|
0 PyTorch - .pt True True
|
|
1 TorchScript torchscript .torchscript True True
|
|
2 ONNX onnx .onnx True True
|
|
3 OpenVINO openvino _openvino_model True False
|
|
4 TensorRT engine .engine False True
|
|
5 CoreML coreml .mlmodel True False
|
|
6 TensorFlow SavedModel saved_model _saved_model True True
|
|
7 TensorFlow GraphDef pb .pb True True
|
|
8 TensorFlow Lite tflite .tflite True False
|
|
9 TensorFlow Edge TPU edgetpu _edgetpu.tflite False False
|
|
10 TensorFlow.js tfjs _web_model False False
|
|
11 PaddlePaddle paddle _paddle_model True True
|
|
"""
|
|
from ultralytics.yolo.engine.exporter import export_formats
|
|
|
|
print(export_formats())
|
|
|
|
model = YOLO("yolov8n.yaml")
|
|
model.export(format='torchscript')
|
|
model.export(format='onnx')
|
|
model.export(format='openvino')
|
|
model.export(format='coreml')
|
|
model.export(format='paddle')
|
|
|
|
|
|
def test():
|
|
test_model_forward()
|
|
test_model_info()
|
|
test_model_fuse()
|
|
test_visualize_preds()
|
|
test_val()
|
|
test_model_resume()
|
|
test_model_train_pretrained()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test()
|