Python load_model Examples

Example #1

File: test_load_model.py Project: toslunar/pytorch-pfn-extras

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"))

Example #2

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"))