def test_yolov5m_r40(self):
images_one, images_two = self.get_test_images()
images_dummy = [torch.ones(3, 100, 100) * 0.3]
model = yolov5m(upstream_version='v4.0',
export_friendly=True,
pretrained=True)
model.eval()
model(images_one)
# Test exported model on images of different size, or dummy input
self.run_model(model, [(images_one, ), (images_two, ),
(images_dummy, )],
input_names=["images_tensors"],
output_names=["outputs"],
dynamic_axes={
"images_tensors": [0, 1, 2],
"outputs": [0, 1, 2]
},
tolerate_small_mismatch=True)
# Test exported model for an image with no detections on other images
self.run_model(model, [(images_dummy, ), (images_one, )],
input_names=["images_tensors"],
output_names=["outputs"],
dynamic_axes={
"images_tensors": [0, 1, 2],
"outputs": [0, 1, 2]
},
tolerate_small_mismatch=True)
def test_yolov5m_script(self):
model = yolov5m(pretrained=True)
model.eval()
scripted_model = torch.jit.script(model)
scripted_model.eval()
x = [torch.rand(3, 416, 320), torch.rand(3, 480, 352)]
out = model(x)
out_script = scripted_model(x)
self.assertTrue(out[0]["scores"].equal(out_script[0]["scores"]))
self.assertTrue(out[0]["labels"].equal(out_script[0]["labels"]))
self.assertTrue(out[0]["boxes"].equal(out_script[0]["boxes"]))