(name, dtype, input_shape=(1, 3, 224, 224))
| 60 | |
| 61 | |
| 62 | def get_network(name, dtype, input_shape=(1, 3, 224, 224)): |
| 63 | def download_model(model_url, name): |
| 64 | model_path = download_testdata(model_url, name + ".onnx", module="onnx") |
| 65 | onnx_model = onnx.load(model_path) |
| 66 | |
| 67 | shape_dict = {"x": input_shape} |
| 68 | mod = from_onnx(onnx_model, shape_dict) |
| 69 | return mod |
| 70 | |
| 71 | def create_model(name): |
| 72 | if "vgg11" == name: |
| 73 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/vgg11_Opset18_timm/vgg11_Opset18.onnx" |
| 74 | elif "mobilenetv3" == name: |
| 75 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/mobilenetv3_large_100_miil_Opset17_timm/mobilenetv3_large_100_miil_Opset17.onnx" |
| 76 | elif "alexnet" == name: |
| 77 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/alexnet_Opset17_torch_hub/alexnet_Opset17.onnx" |
| 78 | elif "resnet50" == name: |
| 79 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/resnet50_Opset18_timm/resnet50_Opset18.onnx" |
| 80 | elif "resnet34" == name: |
| 81 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/resnet34_Opset18_timm/resnet34_Opset18.onnx" |
| 82 | elif "resnet18" == name: |
| 83 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/resnet18_Opset18_timm/resnet18_Opset18.onnx" |
| 84 | elif "squeezenet" == name: |
| 85 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/squeezenet1_1_Opset18_torch_hub/squeezenet1_1_Opset18.onnx" |
| 86 | elif "vgg16" == name: |
| 87 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/vgg16_Opset18_timm/vgg16_Opset18.onnx" |
| 88 | elif "vgg19" == name: |
| 89 | model_url = "https://github.com/onnx/models/raw/bec48b6a70e5e9042c0badbaafefe4454e072d08/Computer_Vision/vgg19_Opset18_timm/vgg19_Opset18.onnx" |
| 90 | else: |
| 91 | assert False, f"Not supported model {name}" |
| 92 | |
| 93 | return download_model(model_url, name) |
| 94 | |
| 95 | mod = create_model(name) |
| 96 | return mod, {"data": (input_shape, dtype)} |
| 97 | |
| 98 | |
| 99 | @pytest.mark.parametrize( |
no test coverage detected
searching dependent graphs…