mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
cf9cad31df
This PR is part of a series attempting to re-submit https://github.com/pytorch/pytorch/pull/134592 as smaller PRs. Add missing `if __name__ == "__main__":` guards to some tests. Pull Request resolved: https://github.com/pytorch/pytorch/pull/154716 Approved by: https://github.com/Skylion007
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
# Owner(s): ["module: onnx"]
|
|
|
|
import os
|
|
import unittest
|
|
|
|
import onnx_test_common
|
|
import parameterized
|
|
import PIL
|
|
import torchvision
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.testing._internal import common_utils
|
|
|
|
|
|
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())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
common_utils.run_tests()
|