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