def load_bytes(b):
b.seek(0)
x = onnx.load(b)
# doc_string has stack traces - let's remove them to make comparison
# sane
onnx.helper.strip_doc_string(x)
return x
def onnx_to_coreml(onnx_filename, core_ml_filename):
onnx_model = onnx.load(onnx_filename)
coreml_model = convert(onnx_model,
minimum_ios_deployment_target="13",
image_input_names=['image'],
preprocessing_args={
"image_scale": 1 / 255.0,
"red_bias": 0,
"green_bias": 0,
"blue_bias": 0
})
builder = coremltools.models.neural_network.NeuralNetworkBuilder(
spec=coreml_model.get_spec())
save_spec(builder.spec, core_ml_filename)
coremltools.models.neural_network.printer.print_network_spec(builder.spec)
def stylize_onnx_caffe2(content_image, args):
"""
Read ONNX model and run it using Caffe2
"""
assert not args.export_onnx
import onnx
import onnx_caffe2.backend
model = onnx.load(args.model)
prepared_backend = onnx_caffe2.backend.prepare(model, device='CUDA' if args.cuda else 'CPU')
inp = {model.graph.input[0].name: content_image.numpy()}
c2_out = prepared_backend.run(inp)[0]
return torch.from_numpy(c2_out)
def test_duplicated_output_node(self):
class DuplicatedOutputNet(torch.nn.Module):
def __init__(self, input_size, num_classes):
super(DuplicatedOutputNet, self).__init__()
self.fc1 = torch.nn.Linear(input_size, num_classes)
def forward(self, input0, input1):
out1 = self.fc1(input0)
out2 = self.fc1(input1)
return out1, out1, out2, out1, out2
N, D_in, H, D_out = 64, 784, 500, 10
pt_model = DuplicatedOutputNet(D_in, D_out)
f = io.BytesIO()
x = torch.randn(N, D_in)
dynamic_axes = {
'input0': {0: 'input0_dim0', 1: 'input0_dim1'},
'input1': {0: 'input1_dim0', 1: 'input1_dim1'},
'output-0': {0: 'output-0_dim0', 1: 'output-0_dim1'},
'output-1': {0: 'output-1_dim0', 1: 'output-1_dim1'},
'output-2': {0: 'output-2_dim0', 1: 'output-2_dim1'},
'output-3': {0: 'output-3_dim0', 1: 'output-3_dim1'},
'output-4': {0: 'output-4_dim0', 1: 'output-4_dim1'}}
torch.onnx.export(pt_model,
(x, x),
f,
input_names=['input0', 'input1'],
output_names=['output-0', 'output-1', 'output-2', 'output-3', 'output-4'],
do_constant_folding=False,
training=torch.onnx.TrainingMode.TRAINING,
dynamic_axes=dynamic_axes,
verbose=True,
keep_initializers_as_inputs=True)
graph = onnx.load(io.BytesIO(f.getvalue()))
assert graph.graph.input[0].name == "input0"
assert graph.graph.input[1].name == "input1"
for i in range(5):
assert graph.graph.output[i].name == "output-" + str(i)
assert graph.graph.node[0].op_type == "Gemm"
assert graph.graph.node[1].op_type == "Identity"
assert graph.graph.node[2].op_type == "Identity"
assert graph.graph.node[3].op_type == "Gemm"
assert graph.graph.node[4].op_type == "Identity"
def get_model(opts):
path = "{}/{}/".format(opts.model_path, opts.model_name)
log_path = "{}/".format(opts.log_path)
filename = "model.onnx"
if not os.path.exists(path):
print("Creating models directory")
os.makedirs(path)
if not os.path.exists(log_path):
print("Creating logs directory")
os.makedirs(log_path)
dataset = "imagenet"
if opts.url:
# Monkey patch an alternate URL into pretrained models package
pretrainedmodels.models.resnext.pretrained_settings[
opts.model_name][dataset]["url"] = opts.url
print(f"Download URL set to {opts.url}")
pretrained_model_base_path = opts.pretrained_model_path
if not pretrained_model_base_path:
# Get the model. If it doesn't exist it will be downloaded
if not os.path.isfile(path + filename):
print(f"Downloading model to {path + filename}")
pretrained_model_path = path + filename
# Create the right input shape
dummy_input = torch.randn(1, 3, 224, 224)
model = pretrainedmodels.__dict__[opts.model_name](num_classes=1000,
pretrained=dataset)
torch.onnx.export(model, dummy_input, pretrained_model_path)
else:
pretrained_model_path = os.path.join(pretrained_model_base_path,
opts.model_name, filename)
onnx_model = onnx.load(pretrained_model_path)
onnx_model.graph.input[0].type.tensor_type.shape.dim[
0].dim_value = opts.micro_batch_size
print(
f"Converting model to batch size {opts.micro_batch_size} and saving to {path + 'model_' + str(opts.micro_batch_size) + '.onnx'}"
)
onnx.save(onnx_model, path + f"model_{opts.micro_batch_size}.onnx")
def main(args):
import os
from config import Config
total_config = Config()
if not args.dataset or args.dataset not in total_config.DATASETS.keys():
raise Exception("specify one of the datasets to use in {}".format(
list(total_config.DATASETS.keys())))
if not args.snapshot or not os.path.isfile(args.snapshot):
raise Exception("invalid snapshot")
dataset = args.dataset
dataset_class = total_config.DATASETS[dataset]
dataset_params = total_config.DATASET_PARAMS[dataset]
model = YoloNet(dataset_config=dataset_params)
model.load_state_dict(torch.load(args.snapshot)["state_dict"])
model.eval()
if args.batch_size:
batch_size = args.batch_size
else:
batch_size = 1
x = torch.randn(batch_size, 3, dataset_params["img_h"],
dataset_params["img_w"])
torch.onnx.export(
model,
x,
args.onnx_weight_file,
verbose=True,
input_names=["input"],
output_names=["output"],
do_constant_folding=True,
operator_export_type=torch.onnx.OperatorExportTypes.ONNX,
opset_version=11,
)
if args.batch_size:
return
import onnx
mp = onnx.load(args.onnx_weight_file)
mp.graph.input[0].type.tensor_type.shape.dim[0].dim_param = "None"
mp.graph.output[0].type.tensor_type.shape.dim[0].dim_param = "None"
onnx.save(mp, "output.onnx")
def test_duplicated_output_node(self):
class DuplicatedOutputNet(torch.nn.Module):
def __init__(self, input_size, num_classes):
super(DuplicatedOutputNet, self).__init__()
self.fc1 = torch.nn.Linear(input_size, num_classes)
def forward(self, input0, input1):
out1 = self.fc1(input0)
out2 = self.fc1(input1)
return out1, out1, out2, out1, out2
N, D_in, H, D_out = 64, 784, 500, 10
pt_model = DuplicatedOutputNet(D_in, D_out)
f = io.BytesIO()
x = torch.randn(N, D_in)
dynamic_axes = {
"input0": {0: "input0_dim0", 1: "input0_dim1"},
"input1": {0: "input1_dim0", 1: "input1_dim1"},
"output-0": {0: "output-0_dim0", 1: "output-0_dim1"},
"output-1": {0: "output-1_dim0", 1: "output-1_dim1"},
"output-2": {0: "output-2_dim0", 1: "output-2_dim1"},
"output-3": {0: "output-3_dim0", 1: "output-3_dim1"},
"output-4": {0: "output-4_dim0", 1: "output-4_dim1"}}
torch.onnx.export(pt_model,
(x, x),
f,
input_names=["input0", "input1"],
output_names=["output-0", "output-1", "output-2", "output-3", "output-4"],
do_constant_folding=False,
training=torch.onnx.TrainingMode.TRAINING,
dynamic_axes=dynamic_axes,
verbose=True,
keep_initializers_as_inputs=True)
graph = onnx.load(io.BytesIO(f.getvalue()))
self.assertEqual(graph.graph.input[0].name, "input0")
self.assertEqual(graph.graph.input[1].name, "input1")
for i in range(5):
self.assertEqual(graph.graph.output[i].name, f"output-{i}")
self.assertEqual(graph.graph.node[0].op_type, "Gemm")
self.assertEqual(graph.graph.node[1].op_type, "Identity")
self.assertEqual(graph.graph.node[2].op_type, "Identity")
self.assertEqual(graph.graph.node[3].op_type, "Gemm")
self.assertEqual(graph.graph.node[4].op_type, "Identity")
def stylize_onnx_caffe2(content_image, args):
"""
Read ONNX model and run it using Caffe2
"""
assert not args.export_onnx
import onnx
import onnx_caffe2.backend
model = onnx.load(args.model)
prepared_backend = onnx_caffe2.backend.prepare(model, device='CUDA' if args.cuda else 'CPU')
inp = {model.graph.input[0].name: content_image.numpy()}
c2_out = prepared_backend.run(inp)[0]
return torch.from_numpy(c2_out)
def quantize_onnx(onnx_path, quantized_onnx_path):
"""
Quantize the weights of the model from float32 to in8 to allow very efficient inference on modern CPU.
"""
import onnx
from onnxruntime.quantization import QuantizationMode, quantize
onnx_model = onnx.load(onnx_path)
quantized_model = quantize(
model=onnx_model,
quantization_mode=QuantizationMode.IntegerOps,
force_fusions=True,
symmetric_weight=True,
)
# Save model
onnx.save_model(quantized_model, quantized_onnx_path)
def check_onnx_model(onnx_path):
# Load the ONNX model
model = onnx.load(
onnx_path
) # load the saved model and will output a onnx.ModelProto structure
print('load', onnx_path)
# Check that the IR is well formed
try:
onnx.checker.check_model(
model
) # verify the model’s structure and confirm that the model has a valid schema.
print('check pass!')
finally: # except
print('check error!')
# Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph) # 没有输出? 因为 check 失败了?
def test_custom_opsets_gelu(self):
def gelu(g, self):
return g.op("com.microsoft::Gelu", self).setType(self.type())
register_custom_op_symbolic("::gelu", gelu, 1)
model = torch.nn.GELU()
x = torch.randn(3, 3)
f = io.BytesIO()
torch.onnx.export(model, (x, ),
f,
opset_version=self.opset_version,
custom_opsets={"com.microsoft": 1})
graph = onnx.load(io.BytesIO(f.getvalue()))
assert graph.graph.node[0].op_type == "Gelu"
assert graph.opset_import[0].version == self.opset_version
assert graph.opset_import[1].domain == 'com.microsoft'
assert graph.opset_import[1].version == 1
def check_model(model_export_path, x, torch_out):
model = onnx.load(model_export_path)
prepared_backend = onnx_caffe2_backend.prepare(model)
W = {model.graph.input[0].name: x.data.cpu().numpy()}
c2_out = prepared_backend.run(W)[0]
np.testing.assert_almost_equal(torch_out.data.cpu().numpy(),
c2_out,
decimal=3)
print(
"Exported model has been executed on Caffe2 backend, and the result looks good!"
)
return prepared_backend
def test_custom_opsets_gelu(self):
self.addCleanup(unregister_custom_op_symbolic, "::gelu", 1)
def gelu(g, self):
return g.op("com.microsoft::Gelu", self).setType(self.type())
register_custom_op_symbolic("::gelu", gelu, 1)
model = torch.nn.GELU()
x = torch.randn(3, 3)
f = io.BytesIO()
torch.onnx.export(model, (x, ), f,
opset_version=self.opset_version, custom_opsets={"com.microsoft": 1})
graph = onnx.load(io.BytesIO(f.getvalue()))
self.assertEqual(graph.graph.node[0].op_type, "Gelu")
self.assertEqual(graph.opset_import[0].version, self.opset_version)
self.assertEqual(graph.opset_import[1].domain, "com.microsoft")
self.assertEqual(graph.opset_import[1].version, 1)
def export_to_onnx(self, model, input, input_names):
traced = torch.jit.trace(model, input)
buf = io.BytesIO()
torch.jit.save(traced, buf)
buf.seek(0)
model = torch.jit.load(buf)
f = io.BytesIO()
torch.onnx.export(model,
input,
f,
input_names=input_names,
operator_export_type=torch.onnx.OperatorExportTypes.
ONNX_ATEN_FALLBACK)
f.seek(0)
onnx_model = onnx.load(f)
return onnx_model
def check_onnx_opsets_operator(module,
x,
ops,
opset_versions,
training=torch.onnx.TrainingMode.EVAL,
input_names=None,
dynamic_axes=None):
for opset_version in opset_versions:
f = io.BytesIO()
torch.onnx.export(module,
x,
f,
opset_version=opset_version,
training=training,
input_names=input_names,
dynamic_axes=dynamic_axes)
model = onnx.load(io.BytesIO(f.getvalue()))
check_onnx_opset_operator(model, ops[opset_version], opset_version)
def convert_onnx(model, val_loader, vocab):
# exports the pytorch model to onnx
# accepts source and target tensor with batch_size and padded_length of variable sizes
for val_batch in val_loader:
source, target = process_batch(val_batch)
break # only one example is needed
torch.onnx.export(model,
(source, target),
"paraphraser.onnx",
opset_version=11,
input_names = ['input'],
output_names = ['output'],
dynamic_axes={'input' : {0 : 'padded_size', 1: 'padded_size'},
'output' : {0 : 'batch_size'}})
onnx_model = onnx.load("paraphraser.onnx")
onnx.checker.check_model(onnx_model)
def convert_and_validate_to_onnx(model, model_name, sz, input_names,
output_names):
dummy_input = Variable(torch.randn(3, sz, sz)).cuda()
torch.onnx.export(model, dummy_input, \
model_name, input_names=input_names,
output_names=output_names, verbose=True)
# Check again by onnx
# Load the ONNX model
onnx_model = onnx.load(model_name)
# Check that the IR is well formed
onnx.checker.check_model(onnx_model)
# Print a human readable representation of the graph
# onnx.helper.printable_graph(onnx_model.graph)
return onnx_model
def pytorch_to_onnx(torch_model, dummy_input, model_name):
#dummy_input = Variable(torch.rand(1, 1, 96, 96))
torch_model.cpu()
torch_model.train(False)
torch_model.convert_to_onnx = True
torch.onnx.export(torch_model,
dummy_input,
f"trained_models/{model_name}.proto",
verbose=False,
export_params=True)
# test if it works
onnx_model = onnx.load(f"trained_models/{model_name}.proto")
onnx.checker.check_model(onnx_model)
print(onnx.helper.printable_graph(onnx_model.graph))
rep = backend.prepare(onnx_model, device="CPU")
print(rep.run(dummy_input.data.cpu().numpy())[0])
return rep
def test_resnet_to_coreml():
onnx_model = onnx.load(resnet_onnx_filename)
onnx.checker.check_model(onnx_model)
from onnx_coreml import convert
coreml_model = convert(onnx_model,
minimum_ios_deployment_target="13",
image_input_names=['image'],
preprocessing_args={
"image_scale": 1 / 255.0,
"red_bias": 0,
"green_bias": 0,
"blue_bias": 0
})
yolo_model = coremltools.models.MLModel(coreml_model.get_spec())
yolo_model.save(resnet_coreml_filename)
def export_onnx(model, size, filename):
model_aug = nn.Sequential(transforms.Normalize(), model).cpu()
dummy = torch.ByteTensor(1, int(size[1]), int(size[0]), 3)
torch.onnx.export(
model_aug, # model being run
dummy, # model input (or a tuple for multiple inputs)
filename, # where to save the model (can be a file or file-like object)
export_params=
True, # store the trained parameter weights inside the model file
# opset_version=10, # the ONNX version to export the model to
do_constant_folding=
True, # wether to execute constant folding for optimization
input_names=['input'], # the model's input names
output_names=['location',
'classification'], # the model's output names
dynamic_axes={})
onnx_model = onnx.load(filename)
onnx.checker.check_model(onnx_model)
def generic_test(self, model, sample_inputs, input_names=None):
pt_inputs = tuple(torch.from_numpy(x) for x in sample_inputs)
model.qconfig = torch.quantization.default_qconfig
q_model = torch.quantization.prepare(model, inplace=False)
q_model = torch.quantization.convert(q_model, inplace=False)
pytorch_res = q_model(*pt_inputs)
f = io.BytesIO()
torch.onnx.export(q_model,
pt_inputs,
f,
input_names=input_names,
operator_export_type=torch.onnx.OperatorExportTypes.
ONNX_ATEN_FALLBACK)
f.seek(0)
onnx_model = onnx.load(f)
caffe_res = c2.run_model(onnx_model,
dict(zip(input_names, sample_inputs)))[0]
np.testing.assert_almost_equal(pytorch_res.numpy(),
caffe_res,
decimal=3)
def test_custom_opsets_inverse(self):
class CustomInverse(torch.nn.Module):
def forward(self, x):
return torch.inverse(x) + x
def inverse(g, self):
return g.op("com.microsoft::Inverse", self).setType(self.type())
register_custom_op_symbolic("::inverse", inverse, 1)
model = CustomInverse()
x = torch.randn(2, 3, 3)
f = io.BytesIO()
torch.onnx.export(model, (x, ), f,
opset_version=self.opset_version, custom_opsets={"com.microsoft": 1})
graph = onnx.load(io.BytesIO(f.getvalue()))
self.assertEqual(graph.graph.node[0].op_type, "Inverse")
self.assertEqual(graph.opset_import[0].version, self.opset_version)
self.assertEqual(graph.opset_import[1].domain, "com.microsoft")
self.assertEqual(graph.opset_import[1].version, 1)
def generic_test(self, model, sample_inputs, input_names=None, permute=False):
pt_inputs = tuple(torch.from_numpy(x) for x in sample_inputs)
model.qconfig = torch.quantization.default_qconfig
q_model = torch.quantization.prepare(model, inplace=False)
q_model = torch.quantization.convert(q_model, inplace=False)
if permute:
# Permute input to caffe2 to be NHWC layout
X_nhwc = np.ascontiguousarray(sample_inputs[0].transpose([0, 2, 3, 1]))
pytorch_res = q_model(*pt_inputs)
f = io.BytesIO()
torch.onnx.export(q_model, pt_inputs, f, input_names=input_names, operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
f.seek(0)
onnx_model = onnx.load(f)
if permute:
caffe_res = c2.run_model(onnx_model, dict(zip(input_names, (X_nhwc,))))[0]
np.testing.assert_almost_equal(pytorch_res.permute(0, 2, 3, 1).numpy(), caffe_res, decimal=3)
else:
caffe_res = c2.run_model(onnx_model, dict(zip(input_names, sample_inputs)))[0]
np.testing.assert_almost_equal(pytorch_res.numpy(), caffe_res, decimal=3)
def convert_to_onnx(pytorch_model, input_shape, output_file, input_names,
output_names):
"""Convert PyTorch model to ONNX and check the resulting onnx model"""
pytorch_model.eval()
onnx_input_shape = torch.randn(input_shape)
torch.onnx.export(pytorch_model,
onnx_input_shape,
output_file,
verbose=True,
input_names=input_names,
output_names=output_names)
# Model check after conversion
import onnx
model_from_onnx = onnx.load(output_file)
try:
onnx.checker.check_model(model_from_onnx)
print('ONNX check passed successfully.')
except onnx.onnx_cpp2py_export.checker.ValidationError as exc:
sys.exit('ONNX check failed with error: ' + str(exc))
def onnx_runtime(infer_tensor, onnx_path, torch_output):
onnx_model = onnx.load(onnx_path)
print(f"onnx model type: {type(onnx_model)}")
onnx.checker.check_model(onnx_model)
ort_session = onnxruntime.InferenceSession(onnx_path)
# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(infer_tensor)}
ort_outs = ort_session.run(None, ort_inputs)
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(to_numpy(torch_output),
ort_outs[0],
rtol=1e-03,
atol=1e-05)
print(
"Exported model has been tested with ONNXRuntime, and the result looks good!"
)
return ort_outs[0]
def export_to_onnx(self, model, input, input_names):
traced = torch.jit.trace(model, input)
buf = io.BytesIO()
torch.jit.save(traced, buf)
buf.seek(0)
model = torch.jit.load(buf)
f = io.BytesIO()
torch.onnx.export(
model,
input,
f,
input_names=input_names,
operator_export_type=torch.onnx.OperatorExportTypes.
ONNX_ATEN_FALLBACK,
# Caffe2 doesn't support newer opset versions
opset_version=9)
f.seek(0)
onnx_model = onnx.load(f)
return onnx_model
def test_conv_variable_length(self):
x = torch.ones(5, 3, 6, 6, requires_grad=True)
model = torch.nn.Conv2d(3, 2, 3)
dynamic_axes = {"input_1": [0, 2, 3], "output_1": {0: "output_1_variable_dim_0", 1: "output_1_variable_dim_1"}}
model_proto_file = tempfile.NamedTemporaryFile()
torch.onnx.export(model, x, model_proto_file.name, verbose=True, input_names=["input_1"], output_names=["output_1"],
dynamic_axes=dynamic_axes)
import onnx
onnx_model = onnx.load(model_proto_file.name)
onnx.checker.check_model(onnx_model)
# Asserting the default dynamic axes names are generated when custom names are not provided
assert(onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_param == "input_1_dynamic_axes_1")
assert(onnx_model.graph.input[0].type.tensor_type.shape.dim[2].dim_param == "input_1_dynamic_axes_2")
assert(onnx_model.graph.input[0].type.tensor_type.shape.dim[3].dim_param == "input_1_dynamic_axes_3")
# Asserting the custom names are applied when provided
assert(onnx_model.graph.output[0].type.tensor_type.shape.dim[0].dim_param == "output_1_variable_dim_0")
assert(onnx_model.graph.output[0].type.tensor_type.shape.dim[1].dim_param == "output_1_variable_dim_1")
def test_conv_variable_length(self):
x = torch.ones(5, 3, 6, 6, requires_grad=True)
model = torch.nn.Conv2d(3, 2, 3)
y = model(x)
dynamic_axes = {
'input_1': [0, 2, 3],
'output_1': {
0: 'output_1_variable_dim_0',
1: 'output_1_variable_dim_1'
}
}
model_proto_name = 'conv2d.onnx'
torch.onnx.export(model,
x,
model_proto_name,
verbose=True,
input_names=["input_1"],
output_names=["output_1"],
example_outputs=y,
dynamic_axes=dynamic_axes)
import onnx
onnx_model = onnx.load(model_proto_name)
onnx.checker.check_model(onnx_model)
# Asserting the default dynamic axes names are generated when custom names are not provided
assert (onnx_model.graph.input[0].type.tensor_type.shape.dim[0].
dim_param == "input_1_dynamic_axes_1")
assert (onnx_model.graph.input[0].type.tensor_type.shape.dim[2].
dim_param == "input_1_dynamic_axes_2")
assert (onnx_model.graph.input[0].type.tensor_type.shape.dim[3].
dim_param == "input_1_dynamic_axes_3")
# Asserting the custom names are applied when provided
assert (onnx_model.graph.output[0].type.tensor_type.shape.dim[0].
dim_param == "output_1_variable_dim_0")
assert (onnx_model.graph.output[0].type.tensor_type.shape.dim[1].
dim_param == "output_1_variable_dim_1")
def export_onnx_and_check(model,
model_name: str = 'model.onnx',
sz: int = 224,
input_names: List[str] = None,
output_names: List[str] = None):
input_names = input_names or ['input']
output_names = output_names or ['output']
x = torch.randn(3, sz, sz)
dummy_input = Variable(x).cuda()
torch.onnx.export(model,
dummy_input,
model_name,
input_names=input_names,
output_names=output_names,
verbose=True)
onnx_model = onnx.load(model_name)
onnx.checker.check_model(onnx_model)
return onnx_model
def load_model_from_ONNX(self, modelName="yolo2"):
"""
import (.onnx) file
build model (model contains weights and structure)
"""
import onnx
# Load the ONNX model
onnxfilepath = "./onnx/{}.onnx".format(modelName)
if os.path.isfile(os.getcwd() + onnxfilepath[1:]):
print('Load onnx file [{}]...start'.format(onnxfilepath))
st = time.time()
model = onnx.load(onnxfilepath)
print('Load onnx file [{}]...done, {:.2f} sec'.format(
onnxfilepath,
time.time() - st))
# Check that the IR is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
model_flat_IR = onnx.helper.printable_graph(model.graph)
return model, model_flat_IR
else:
print('Onnx file path Error!')
def quantize_onnx(onnx_path, quantized_onnx_path):
"""
Quantize the weights of the model from float32 to in8 to allow very efficient inference on modern CPU.
"""
import onnx
from onnxruntime.quantization import QuantizationMode, quantize
onnx_model = onnx.load(onnx_path)
# Discussed with @yufenglee from ONNX runtime, this will be address in the next release of onnxruntime
print(
"As of onnxruntime 1.4.0, models larger than 2GB will fail to quantize due to protobuf constraint.\n"
"This limitation will be removed in the next release of onnxruntime.")
quantized_model = quantize(
model=onnx_model,
quantization_mode=QuantizationMode.IntegerOps,
force_fusions=True,
symmetric_weight=True,
)
# Save model
onnx.save_model(quantized_model, quantized_onnx_path)
def onnx_inference(args):
# Load the ONNX model
model = onnx.load("models/deepspeech_{}.onnx".format(args.continue_from))
# Check that the IR is well formed
onnx.checker.check_model(model)
onnx.helper.printable_graph(model.graph)
print("model checked, preparing backend!")
rep = backend.prepare(model, device="CPU") # or "CPU"
print("running inference!")
# Hard coded input dim
inputs = np.random.randn(16, 1, 161, 129).astype(np.float32)
start = time.time()
outputs = rep.run(inputs)
print("time used: {}".format(time.time() - start))
# To run networks with more than one input, pass a tuple
# rather than a single numpy ndarray.
print(outputs[0])
######################################################################
# ``torch_out`` is the output after executing the model. Normally you can
# ignore this output, but here we will use it to verify that the model we
# exported computes the same values when run in Caffe2.
#
# Now let's take the ONNX representation and use it in Caffe2. This part
# can normally be done in a separate process or on another machine, but we
# will continue in the same process so that we can verify that Caffe2 and
# PyTorch are computing the same value for the network:
#
import onnx
import onnx_caffe2.backend
# Load the ONNX ModelProto object. model is a standard Python protobuf object
model = onnx.load("super_resolution.onnx")
# prepare the caffe2 backend for executing the model this converts the ONNX model into a
# Caffe2 NetDef that can execute it. Other ONNX backends, like one for CNTK will be
# availiable soon.
prepared_backend = onnx_caffe2.backend.prepare(model)
# run the model in Caffe2
# Construct a map from input names to Tensor data.
# The graph of the model itself contains inputs for all weight parameters, after the input image.
# Since the weights are already embedded, we just need to pass the input image.
# Set the first input.
W = {model.graph.input[0].name: x.data.numpy()}
# Run the Caffe2 net:
# you can combine your ATen ops with standard onnx ones
x = nn.ReLU()(x)
return MyFunction.apply(x, y)
torch.onnx.export(MyModule(),
(Variable(torch.ones(3,4)), Variable(torch.ones(3,4))),
"output.onnx",
verbose=True)
# prints the graph for debugging:
# graph(%1 : Float(3, 4)
# %2 : Float(3, 4)) {
# %3 : Float(3, 4) = Relu(%1), uses = [%4.i0, %4.i1];
# %4 : UNKNOWN_TYPE = ATen[operator=mul](%3, %3), uses = [%5.i0];
# %5 : Float(3, 4) = ATen[operator=add](%4, %2), uses = [%0.i0];
# return (%5);
# }
graph = onnx.load("output.onnx")
a = np.random.randn(3, 4).astype(np.float32)
b = np.random.randn(3, 4).astype(np.float32)
prepared_backend = onnx_caffe2.backend.prepare(graph)
W = {graph.graph.input[0].name: a, graph.graph.input[1].name: b}
c2_out = prepared_backend.run(W)[0]
x = np.maximum(a, 0)
r = x*x + b
np.testing.assert_array_almost_equal(r, c2_out)