def test_stripped_onnx_load_model():
model = Net()
outdir = "out/stripped_load_model_test"
tou.export_testcase(model, torch.rand(1, 1, 28, 28), outdir,
strip_large_tensor_data=True, training=True,
do_constant_folding=False)
tou.load_model(os.path.join(outdir, "model.onnx"))
def _helper(model, args, d, **kwargs):
output_dir = _get_output_dir(d)
if 'training' not in kwargs:
kwargs['training'] = model.training
if 'do_constant_folding' not in kwargs:
kwargs['do_constant_folding'] = False
export_testcase(model, args, output_dir, **kwargs)
return output_dir
def test_onnx_load_model():
model = Net()
outdir = "out/load_model_test"
tou.export_testcase(model,
torch.rand(1, 1, 28, 28),
outdir,
training=True,
do_constant_folding=False)
tou.load_model(os.path.join(outdir, "model.onnx"))
def _helper(model, args, d, use_pfto=True, **kwargs):
output_dir = _get_output_dir(d)
if 'training' not in kwargs:
kwargs['training'] = model.training
if 'do_constant_folding' not in kwargs:
kwargs['do_constant_folding'] = False
if 'metadata' not in kwargs:
kwargs["metadata"] = False
export_testcase(model, args, output_dir, use_pfto=use_pfto, **kwargs)
return output_dir
def test_export_testcase_return_output():
model = nn.Sequential(nn.Linear(5, 10, bias=False))
x = torch.zeros((2, 5))
output_dir = _get_output_dir('export_filename')
if pytorch_pfn_extras.requires("1.6.0"):
with pytest.warns(UserWarning):
(out, ) = export_testcase(model, x, output_dir, return_output=True)
else:
(out, ) = export_testcase(model, x, output_dir, return_output=True)
assert os.path.isfile(os.path.join(output_dir, 'model.onnx'))
expected_out = torch.zeros((2, 10)) # check only shape size
np.testing.assert_allclose(out.detach().cpu().numpy(),
expected_out.detach().cpu().numpy())
def _helper(model, args, d, **kwargs):
output_dir = _get_output_dir(d)
if 'training' not in kwargs:
kwargs['training'] = model.training
export_testcase(model, args, output_dir, **kwargs)
return output_dir