Function test_onnx_save

tests/test_utils.py:809–830

Test the ability to save `net` in ONNX format, reload it and validate with runtime. The value `inputs` is forward-passed through the `net` without gradient accumulation to do onnx export and PyTorch inference. PyTorch model inference is performed with CUDA if available, else CPU.

(net, *inputs, device=None, rtol=1e-4, atol=0.0)

def test_onnx_save(net, *inputs, device=None, rtol=1e-4, atol=0.0):
    """
    Test the ability to save `net` in ONNX format, reload it and validate with runtime.
    The value `inputs` is forward-passed through the `net` without gradient accumulation
    to do onnx export and PyTorch inference.
    PyTorch model inference is performed with CUDA if available, else CPU.
    Saved ONNX model is validated with onnxruntime, if available, else ONNX native implementation.
    """
    # TODO: would be nice to use GPU if available, but it currently causes CI failures.
    device = "cpu"
    _, has_onnxruntime = optional_import("onnxruntime")
    with tempfile.TemporaryDirectory() as tempdir:
        convert_to_onnx(
            model=net,
            filename=os.path.join(tempdir, "model.onnx"),
            verify=True,
            inputs=inputs,
            device=device,
            use_ort=has_onnxruntime,
            rtol=rtol,
            atol=atol,
        )


def download_url_or_skip_test(*args, **kwargs):