1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
# Owner(s): ["module: onnx"]
import os
import unittest
import onnx_test_common
import parameterized
import PIL
import torch
import torchvision
from torch import nn
def _get_test_image_tensor():
data_dir = os.path.join(os.path.dirname(__file__), "assets")
img_path = os.path.join(data_dir, "grace_hopper_517x606.jpg")
input_image = PIL.Image.open(img_path)
# Based on example from https://pytorch.org/hub/pytorch_vision_resnet/
preprocess = torchvision.transforms.Compose(
[
torchvision.transforms.Resize(256),
torchvision.transforms.CenterCrop(224),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
),
]
)
return preprocess(input_image).unsqueeze(0)
# Due to precision error from quantization, check only that the top prediction matches.
class _TopPredictor(nn.Module):
def __init__(self, base_model):
super().__init__()
self.base_model = base_model
def forward(self, x):
x = self.base_model(x)
_, topk_id = torch.topk(x[0], 1)
return topk_id
# TODO: All torchvision quantized model test can be written as single parameterized test case,
# after per-parameter test decoration is supported via #79979, or after they are all enabled,
# whichever is first.
@parameterized.parameterized_class(
("is_script",),
[(True,), (False,)],
class_name_func=onnx_test_common.parameterize_class_name,
)
class TestQuantizedModelsONNXRuntime(onnx_test_common._TestONNXRuntime):
def run_test(self, model, inputs, *args, **kwargs):
model = _TopPredictor(model)
return super().run_test(model, inputs, *args, **kwargs)
def test_mobilenet_v3(self):
model = torchvision.models.quantization.mobilenet_v3_large(
pretrained=True, quantize=True
)
self.run_test(model, _get_test_image_tensor())
@unittest.skip("quantized::cat not supported")
def test_inception_v3(self):
model = torchvision.models.quantization.inception_v3(
pretrained=True, quantize=True
)
self.run_test(model, _get_test_image_tensor())
@unittest.skip("quantized::cat not supported")
def test_googlenet(self):
model = torchvision.models.quantization.googlenet(
pretrained=True, quantize=True
)
self.run_test(model, _get_test_image_tensor())
@unittest.skip("quantized::cat not supported")
def test_shufflenet_v2_x0_5(self):
model = torchvision.models.quantization.shufflenet_v2_x0_5(
pretrained=True, quantize=True
)
self.run_test(model, _get_test_image_tensor())
def test_resnet18(self):
model = torchvision.models.quantization.resnet18(pretrained=True, quantize=True)
self.run_test(model, _get_test_image_tensor())
def test_resnet50(self):
model = torchvision.models.quantization.resnet50(pretrained=True, quantize=True)
self.run_test(model, _get_test_image_tensor())
def test_resnext101_32x8d(self):
model = torchvision.models.quantization.resnext101_32x8d(
pretrained=True, quantize=True
)
self.run_test(model, _get_test_image_tensor())
|
