def import_model(model_file):
"""Imports the ONNX model file, passed as a parameter, into MXNet symbol and parameters.
Operator support and coverage - https://cwiki.apache.org/confluence/display/MXNET/ONNX
Parameters
----------
model_file : str
ONNX model file name
Returns
-------
sym : :class:`~mxnet.symbol.Symbol`
MXNet symbol object
arg_params : dict of ``str`` to :class:`~mxnet.ndarray.NDArray`
Dict of converted parameters stored in ``mxnet.ndarray.NDArray`` format
aux_params : dict of ``str`` to :class:`~mxnet.ndarray.NDArray`
Dict of converted parameters stored in ``mxnet.ndarray.NDArray`` format
"""
graph = GraphProto()
try:
import onnx
except ImportError:
raise ImportError("Onnx and protobuf need to be installed. "
+ "Instructions to install - https://github.com/onnx/onnx")
# loads model file and returns ONNX protobuf object
model_proto = onnx.load(model_file)
sym, arg_params, aux_params = graph.from_onnx(model_proto.graph)
return sym, arg_params, aux_params
def testImportFromPytorch(self):
"""
Test importing a pre-trained pytorch model into tensorflow for inference
"""
# Export trained pytorch model to onnx
pt_model = PytorchMNIST(N)
self._train_pytorch(pt_model)
dummy_input = torch.autograd.Variable(torch.FloatTensor(1, 28 * 28))
torch.onnx.export(pt_model, dummy_input,
os.path.join(self.tmpDir, "torch.onnx"),
input_names=["input"] + pt_model.state_dict().keys(),
output_names=["output"])
_, expected = self._test_pytorch(pt_model)
# Load ONNX model
model = onnx.load(os.path.join(self.tmpDir, "torch.onnx"))
# HACK: Convert model's input shape to dynamic, ignoring batch size
model.graph.input[0].type.tensor_type.shape.dim[0].dim_param = '?'
onnx.checker.check_model(model)
tf_model = onnx_tf.backend.prepare(model)
# Load MNIST test data
(_, _), (x_test, y_test) = keras.datasets.mnist.load_data()
x_test = x_test.reshape(-1, 28 * 28) / 255.0
x_test = (x_test - 0.1307) / 0.3081
with self.test_session(config=CONFIG):
# Compute accuracy
output = tf_model.run(x_test)
predictions = output[0].argmax(axis=1)
accuracy = tf.reduce_mean(tf.cast(predictions == y_test, tf.float32))
self.assertAlmostEqual(accuracy.eval(), expected, places=4)
def get_model_metadata(model_file):
"""
Returns the name and shape information of input and output tensors of the given ONNX model file.
Parameters
----------
model_file : str
ONNX model file name
Returns
-------
model_metadata : dict
A dictionary object mapping various metadata to its corresponding value.
The dictionary will have the following template.
{
'input_tensor_data' : <list of tuples representing the shape of the input paramters>,
'output_tensor_data' : <list of tuples representing the shape of the output
of the model>
}
"""
graph = GraphProto()
try:
import onnx
except ImportError:
raise ImportError("Onnx and protobuf need to be installed. "
+ "Instructions to install - https://github.com/onnx/onnx")
model_proto = onnx.load(model_file)
metadata = graph.get_graph_metadata(model_proto.graph)
return metadata
def import_to_gluon(model_file, ctx):
"""
Imports the ONNX model files, passed as a parameter, into Gluon SymbolBlock object.
Parameters
----------
model_file : str
ONNX model file name
ctx : Context or list of Context
Loads the model into one or many context(s).
Returns
-------
sym_block : :class:`~mxnet.gluon.SymbolBlock`
A SymbolBlock object representing the given model file.
"""
graph = GraphProto()
try:
import onnx
except ImportError:
raise ImportError("Onnx and protobuf need to be installed. Instructions to"
+ " install - https://github.com/onnx/onnx#installation")
model_proto = onnx.load(model_file)
net = graph.graph_to_gluon(model_proto.graph, ctx)
return net
def check_model(): # type: () -> None
parser = argparse.ArgumentParser('check-model')
parser.add_argument('model_pb', type=argparse.FileType('rb'))
args = parser.parse_args()
model = load(args.model_pb)
checker.check_model(model)
def test_export_to_buffer(self):
model = Model()
buffer = model.export_to_buffer()
protobuf_model = onnx.load(BytesIO(buffer.getvalue()))
self.assertEqual(4, len(protobuf_model.graph.input)) # 2 inputs + 2 params
self.assertEqual(4, len(protobuf_model.graph.output))
self.assertEqual("state:float_features", protobuf_model.graph.input[0].name)
self.assertEqual("action:float_features", protobuf_model.graph.input[1].name)
def main():
xmodel = onnx.load(sys.argv[1])
subgraphs = find_subgraphs(xmodel.graph)
if len(sys.argv) == 2:
for name, graph in subgraphs.items():
print('name=%s node=%d' % (name, len(graph.node)))
else:
g = subgraphs[sys.argv[2]]
m = onnx.helper.make_model(g)
onnx.save(m, 'out.onnx')
def compare_graph(onnx_file, nnvm_sym, ishape):
onnx_graph = onnx.load(onnx_file)
onnx_sym, params = nnvm.frontend.from_onnx(onnx_graph)
g1 = nnvm.graph.create(onnx_sym)
g2 = nnvm.graph.create(nnvm_sym)
ishapes = {'input_0': ishape}
graph_attr.set_shape_inputs(g1, ishapes)
graph_attr.set_shape_inputs(g2, ishapes)
g1 = g1.apply("InferShape").apply("SimplifyInference")
g2 = g2.apply("InferShape").apply("SimplifyInference")
graph_util.check_graph_equal(g1, g2)
def convert_onnx(network, net_info, model_root, model_name, copy_params):
model_file = model_root + '/' + model_name + '.onnx'
onnx_model = optimize_onnx(onnx.load(model_file))
onnx_graph = onnx_model.graph
onnx_nodes = onnx_graph.node
onnx_initializer = onnx_graph.initializer
input_infos, output_names = parse_graph(onnx_graph)
net_info['arg']['out_blob_v_s'] = output_names
network.set_net_name(model_name)
network.set_net_arg_dict(net_info['arg'])
convert_input(net_info, network, input_infos)
param_dict = {}
for index, node in enumerate(onnx_nodes):
op_type = node.op_type
if op_type == 'Relu' or op_type == 'PRelu' or op_type == 'Sigmoid':
convert_activate(onnx_nodes, index, onnx_initializer, param_dict, network)
elif op_type == 'BatchNormalization':
convert_batch_norm(onnx_nodes, index, param_dict, network)
elif op_type == 'Add' or op_type == 'Sub' or op_type == 'Mul' or op_type == 'Div':
convert_binary(onnx_nodes, index, network)
elif op_type == 'Concat':
convert_concat(onnx_nodes, index, network)
elif op_type == 'Gemm':
convert_connected(onnx_nodes, index, onnx_initializer, param_dict, network)
elif op_type == 'Conv':
convert_conv(onnx_nodes, index, onnx_initializer, param_dict, network)
elif op_type == 'ConvTranspose':
convert_deconv(onnx_nodes, index, onnx_initializer, param_dict, network)
elif op_type == 'Flatten':
convert_flatten(onnx_nodes, index, network)
elif op_type == 'Transpose':
convert_permute(onnx_nodes, index, network)
elif op_type == 'MaxPool' or op_type == 'AveragePool' or op_type == 'GlobalMaxPool' or op_type == 'GlobalAveragePool':
convert_pooling(onnx_nodes, index, network)
elif op_type == 'Reshape':
convert_reshape(onnx_nodes, index, onnx_initializer, network)
elif op_type == 'Upsample':
convert_resize(onnx_nodes, index, onnx_initializer, network)
elif op_type == 'Softmax':
convert_softmax(onnx_nodes, index, network)
elif op_type == 'Squeeze':
convert_squeeze(onnx_nodes, index, network)
else:
print('Skipping ' + op_type + ', please check!')
if copy_params:
copy_weights(onnx_initializer, param_dict, network)
def prepare_zip_archive(cls, file, device='CPU', **kwargs):
with zipfile.ZipFile(file, mode='r') as z:
with z.open('__MODEL_PROTO', 'r') as f:
model = onnx.load(f);
blob_names = set(z.namelist()) - set('__MODEL_PROTO')
# TODO: make this more efficient
raw_values_dict = {}
for name in blob_names:
with z.open(name, 'r') as blob_file:
raw_values_dict[name] = blob_file.read()
return cls.prepare(model, device, raw_values_dict=raw_values_dict, **kwargs)
def buffer_to_caffe2_netdef(buffer):
"""Creates caffe2 NetDef from buffer object and returns pointer to
input and output blobs and the NetDef."""
protobuf_model = onnx.load(BytesIO(buffer.getvalue()))
input_blob_name = protobuf_model.graph.input[0].name
output_blob_name = protobuf_model.graph.output[0].name
logger.info(
"INPUT BLOB: " + input_blob_name + ". OUTPUT BLOB:" + output_blob_name
)
return (
input_blob_name,
output_blob_name,
caffe2.python.onnx.backend.prepare(protobuf_model),
)
def test_save_no_junk(tmpdir):
""" Test for an issue with save() not having O_TRUNC mode
resulting in possible junk at the end of file if it already exists
"""
try:
import onnx
except ImportError:
pytest.skip('ONNX is not installed.')
filename = os.path.join(str(tmpdir), R'no_junk.onnx')
filename_new = os.path.join(str(tmpdir), R'no_junk_new.onnx')
# Save a large model.
g = C.ceil([1, 2, 3, 4, 5, 6, 7, 8] * 100)
g.save(filename, format=C.ModelFormat.ONNX)
# Save a smaller model using the same file name and a new file name.
g = C.ceil([1, 2, 3, 4, 5, 6, 7] * 100)
g.save(filename, format=C.ModelFormat.ONNX)
g.save(filename_new, format=C.ModelFormat.ONNX)
# Check the files can be loaded as standard ONNX files,
# and both result in the same model.
assert onnx.load(filename) == onnx.load(filename_new)
def main():
args = make_args()
config = configparser.ConfigParser()
utils.load_config(config, args.config)
for cmd in args.modify:
utils.modify_config(config, cmd)
with open(os.path.expanduser(os.path.expandvars(args.logging)), 'r') as f:
logging.config.dictConfig(yaml.load(f))
model_dir = utils.get_model_dir(config)
model = onnx.load(model_dir + '.onnx')
onnx.checker.check_model(model)
init_net, predict_net = onnx_caffe2.backend.Caffe2Backend.onnx_graph_to_caffe2_net(model.graph, device='CPU')
onnx_caffe2.helper.save_caffe2_net(init_net, os.path.join(model_dir, 'init_net.pb'))
onnx_caffe2.helper.save_caffe2_net(predict_net, os.path.join(model_dir, 'predict_net.pb'), output_txt=True)
logging.info(model_dir)
def onnx_input_output_names(onnx_filename):
onnx_model = onnx.load(onnx_filename)
initializer_names = set()
for initializer in onnx_model.graph.initializer:
initializer_names.add(initializer.name)
input_names = []
for input in onnx_model.graph.input:
if input.name not in initializer_names:
input_names.append(input.name)
output_names = []
for output in onnx_model.graph.output:
output_names.append(output.name)
return input_names, output_names
def _roundtrip(self, model_name):
model_dir = Runner(c2)._prepare_model_data(
namedtuple('dummy', ['model_name'])(model_name))
pb_path = os.path.join(model_dir, 'model.pb')
before_roundtrip = onnx.load(pb_path)
with open(pb_path, 'rb') as pb:
after_roundtrip = onnx.load_from_string(pb.read())
assert onnx.helper.printable_graph(before_roundtrip.graph) \
== onnx.helper.printable_graph(after_roundtrip.graph)
with open(pb_path, 'rb') as pb:
assert after_roundtrip.SerializeToString() == pb.read()
def convert_onnx_to_model(onnx_input_path):
"""Reimplementation of the TensorFlow-onnx official tutorial convert the onnx file to specific: model
Parameters
-----------
onnx_input_path : string
the path where you save the onnx file.
References
-----------
- `onnx-tf exporting tutorial <https://github.com/onnx/tutorials/blob/master/tutorials/OnnxTensorflowExport.ipynb>`__
"""
model = onnx.load(onnx_input_path)
tf_rep = prepare(model)
# Image Path
img = np.load("./assets/image.npz")
output = tf_rep.run(img.reshape([1, 784]))
print("The digit is classified as ", np.argmax(output))
def main():
from_dir = sys.argv[1]
to_dir = sys.argv[2]
os.makedirs(to_dir, exist_ok=True)
xmodel = onnx.load(os.path.join(from_dir, 'model.onnx'))
convert_model(xmodel)
onnx.save(xmodel, os.path.join(to_dir, 'model.onnx'))
for test_dir in glob.glob(os.path.join(from_dir, 'test_data_set_*')):
dir_name = os.path.basename(test_dir)
to_test_dir = os.path.join(to_dir, dir_name)
os.makedirs(to_test_dir, exist_ok=True)
for pb_filename in glob.glob(os.path.join(test_dir, '*.pb')):
pb_name = os.path.basename(pb_filename)
tensor = onnx.load_tensor(pb_filename)
convert_tensor(tensor)
onnx.save_tensor(tensor, os.path.join(to_test_dir, pb_name))
def _test_onnx_importer(self, model_name, data_input_index = 0):
model_dir = _download_onnx_model(model_name)
model_def = onnx.load(os.path.join(model_dir, 'model.onnx'))
input_blob_dims = [int(x.dim_value) for x in model_def.graph.input[data_input_index].type.tensor_type.shape.dim]
op_inputs = [x.name for x in model_def.graph.input]
op_outputs = [x.name for x in model_def.graph.output]
print("{}".format(op_inputs))
data = np.random.randn(*input_blob_dims).astype(np.float32)
Y_c2 = c2.run_model(model_def, {op_inputs[data_input_index]: data})
op = convert_onnx_model_to_trt_op(model_def, verbosity=3)
device_option = core.DeviceOption(caffe2_pb2.CUDA, 0)
op.device_option.CopyFrom(device_option)
Y_trt = None
ws = Workspace()
with core.DeviceScope(device_option):
ws.FeedBlob(op_inputs[data_input_index], data)
ws.RunOperatorsOnce([op])
output_values = [ws.FetchBlob(name) for name in op_outputs]
Y_trt = namedtupledict('Outputs', op_outputs)(*output_values)
np.testing.assert_allclose(Y_c2, Y_trt, rtol=1e-3)
def _test(self, modelName):
path = os.path.join(ONNX_MODEL_ZOO_ROOT, modelName, "model.onnx")
try:
o = onnx.load(path)
except FileNotFoundError:
self.skipTest("ONNX model is not found." \
"You may want to download it by running scripts/onnx/download_onnx_model_zoo.py" \
": {}".format(modelName))
onnxInputLayer, = set(map(lambda x: x.name, o.graph.input)) - set(map(lambda x: x.name, o.graph.initializer))
layers = lbann.onnx.o2l.onnxToLbannLayers(
o,
[DATA_LAYER_NAME, LABEL_LAYER_NAME],
{DATA_LAYER_NAME: onnxInputLayer},
)
model = lbann_pb2.Model(layer = layers)
pb = lbann_pb2.LbannPB(model=model)
if SAVE_PROTOTEXT:
with open(os.path.join(DUMP_DIR, "{}.prototext".format(modelName)), "w") as f:
f.write(txtf.MessageToString(pb))
def runBertTrainingTest(gradient_accumulation_steps,
use_mixed_precision,
allreduce_post_accumulation,
use_simple_model_desc=True,
use_internel_loss_scale=False):
model_desc = bert_model_description()
simple_model_desc = remove_extra_info(
model_desc) if use_simple_model_desc else model_desc
learning_rate_description = ort_trainer_learning_rate_description()
device = torch.device("cuda", 0)
torch.manual_seed(1)
onnxruntime.set_seed(1)
onnx_model = onnx.load(get_name("bert_toy_postprocessed.onnx"))
loss_scaler = LossScaler("ort_test_input_loss_scalar",
True) if use_internel_loss_scale else None
model = ORTTrainer(onnx_model,
None,
simple_model_desc,
"LambOptimizer",
map_optimizer_attributes,
learning_rate_description,
device,
postprocess_model=None,
gradient_accumulation_steps=gradient_accumulation_steps,
world_rank=0,
world_size=1,
loss_scaler=loss_scaler,
use_mixed_precision=use_mixed_precision,
allreduce_post_accumulation=allreduce_post_accumulation)
if loss_scaler is None:
loss_scaler = LossScaler(model.loss_scale_input_name, True)
input_ids_batches = []
segment_ids_batches = []
input_mask_batches = []
masked_lm_labels_batches = []
next_sentence_labels_batches = []
batch_size = 16
num_batches = 8
for batch in range(num_batches):
input_ids_batches = [
*input_ids_batches,
generate_sample_batch(model_desc.inputs_[0], batch_size, device)
]
segment_ids_batches = [
*segment_ids_batches,
generate_sample_batch(model_desc.inputs_[1], batch_size, device)
]
input_mask_batches = [
*input_mask_batches,
generate_sample_batch(model_desc.inputs_[2], batch_size, device)
]
masked_lm_labels_batches = [
*masked_lm_labels_batches,
generate_sample_batch(model_desc.inputs_[3], batch_size, device)
]
next_sentence_labels_batches = [
*next_sentence_labels_batches,
generate_sample_batch(model_desc.inputs_[4], batch_size, device)
]
lr_batch_list = [
0.0000000e+00, 4.6012269e-07, 9.2024538e-07, 1.3803681e-06,
1.8404908e-06, 2.3006135e-06, 2.7607362e-06, 3.2208588e-06,
3.6809815e-06
]
actual_losses = []
actual_all_finites = []
for batch_count in range(num_batches):
input_ids = generate_sample_batch(model_desc.inputs_[0], batch_size,
device)
segment_ids = generate_sample_batch(model_desc.inputs_[1], batch_size,
device)
input_mask = generate_sample_batch(model_desc.inputs_[2], batch_size,
device)
masked_lm_labels = generate_sample_batch(model_desc.inputs_[3],
batch_size, device)
next_sentence_labels = generate_sample_batch(model_desc.inputs_[4],
batch_size, device)
lr = lr_batch_list[batch_count]
learning_rate = torch.tensor([lr]).to(device)
training_args = [
input_ids, segment_ids, input_mask, masked_lm_labels,
next_sentence_labels, learning_rate
]
if use_mixed_precision:
if not use_internel_loss_scale:
loss_scale = torch.tensor([loss_scaler.loss_scale_]).to(device)
training_args.append(loss_scale)
actual_loss = model.train_step(*training_args)
if isinstance(actual_loss, (list, tuple)):
assert len(actual_loss) == 2
actual_loss, actual_all_finite = actual_loss
if not use_internel_loss_scale:
loss_scaler.update_loss_scale(actual_all_finite.item())
actual_all_finites = [
*actual_all_finites,
actual_all_finite.cpu().numpy().item(0)
]
actual_losses = [*actual_losses, actual_loss.cpu().numpy().item(0)]
else:
loss = model(*training_args)
actual_losses = [*actual_losses, loss.cpu().numpy().item(0)]
if batch_count == num_batches - 1:
# test eval_step api with fetches at the end of the training.
# if eval_step is called during the training, it will affect the actual training loss (training session is stateful),
eval_loss = model.eval_step(input_ids,
segment_ids,
input_mask,
masked_lm_labels,
next_sentence_labels,
fetches=['loss'])
eval_loss = eval_loss.cpu().numpy().item(0)
# If using internal loss scale, all_finites are handled internally too.
if use_mixed_precision and not use_internel_loss_scale:
return actual_losses, actual_all_finites, eval_loss
else:
return actual_losses, eval_loss
# An example input you would normally provide to your model's forward() method.
input = torch.ones(1, 3, 256, 128)
raw_output = model(input)
print(raw_output.size())
torch.onnx.export(model,
input,
'osnet_x0_25_mcmt17.onnx',
verbose=False,
export_params=True)
print(
"-------------------------check model---------------------------------------\n"
)
try:
onnx_model = onnx.load("osnet_x0_25_mcmt17.onnx")
onnx.checker.check_model(onnx_model)
graph_output = onnx.helper.printable_graph(onnx_model.graph)
with open("graph_output.txt", mode="w") as fout:
fout.write(graph_output)
except:
print("Something went wrong")
import onnxruntime
import numpy as np
ort_session = onnxruntime.InferenceSession("osnet_x0_25_mcmt17.onnx")
def to_numpy(tensor):
return tensor.detach().cpu().numpy(
def __init__(self, onnx_model_path):
self.model = onnx.load(onnx_model_path)
model = Net()
model.eval()
# Convert from PyTorch to ONNX
input_names = ['input']
output_names = ['output']
dummy_input = torch.randn(1, 10, device='cpu')
torch.onnx.export(model,
dummy_input,
'onnx/linear.onnx',
verbose=False,
input_names=input_names,
output_names=output_names)
# Convert From ONNX to TF
onnx_model = onnx.load("onnx/linear.onnx")
tf_rep = prepare(onnx_model, strict=False)
tf_rep.export_graph("onnx/linear.pb")
# Test TF model on random input data.
input_tensor = np.random.random_sample((1, 10))
graph_def = tf.GraphDef()
graph_def.ParseFromString(open('onnx/linear.pb', 'rb').read())
tf.import_graph_def(graph_def, name='')
graph = tf.get_default_graph()
input_node = graph.get_tensor_by_name('input:0')
output_node = graph.get_tensor_by_name('output:0')
with tf.Session() as sess:
args.task_name, )
eval_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
evaluate=True)
# Note that DistributedSampler samples randomly
eval_sampler = SequentialSampler(eval_dataset)
eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, \
batch_size=args.eval_batch_size)
def eval_func(model):
return evaluate_onnxrt(args, model, tokenizer, eval_dataloader)
if args.benchmark:
model = onnx.load(args.model_path)
from lpot.data.datasets.dummy_dataset import DummyDataset
from lpot.data.dataloaders.onnxrt_dataloader import ONNXRTDataLoader
shapes, lows, highs = parse_dummy_input(model, args.benchmark_nums,
args.max_seq_length)
dummy_dataset = DummyDataset(shapes,
low=lows,
high=highs,
dtype="int64")
dummy_dataloader = ONNXRTDataLoader(dummy_dataset)
print(
'---------------------------------------------------------------')
if args.accuracy_only:
results = evaluate_onnxrt(args, model, tokenizer, eval_dataloader)
fontsize=8,
textcoords='offset points',
ha='right',
va='bottom',
bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.5),
arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0'))
###########################################
# Remove one layer at the end
# ---------------------------
#
# The last is often removed before the model is
# inserted in a pipeline. Let's see how to do that.
# First, we need the list of output for every node.
model_onnx = onnx.load(BytesIO(model_bytes))
outputs = []
for node in model_onnx.graph.node:
print(node.name, node.output)
outputs.extend(node.output)
#################################
# We select one of the last one.
selected = outputs[-3]
print("selected", selected)
#################################
# And we tell *OnnxTransformer* to use that
# specific one and to flatten the output
# as the dimension is not a matrix.
def save(self, filename, x, metaData, save_model=False, export_model=False):
"""
Saves the trained model to four files: a JSON file with the settings, a pickled pyTorch state dict
file, and numpy files for the mean and variance of the inputs (used for input scaling).
Also exports model to onnx if export_model is set to True.
Parameters
----------
filename : str
Path to the files. '_settings.json' and '_state_dict.pl' will be added.
save_model : bool, optional
If True, the whole model is saved in addition to the state dict. This is not necessary for loading it
again with Estimator.load(), but can be useful for debugging, for instance to plot the computational graph.
export_model : bool, optional
If True, the whole model is exported to .onnx format to be loaded within a C++ envirnoment.
Returns
-------
None
"""
logger.info("Saving model to %s", filename)
if self.model is None:
raise ValueError("No model -- train or load model before saving!")
# Check paths
create_missing_folders([os.path.dirname(filename)])
# Save settings
logger.debug("Saving settings to %s_settings.json", filename)
settings = self._wrap_settings()
with open(filename + "_settings.json", "w") as f:
json.dump(settings, f)
# Save scaling
if self.x_scaling_stds is not None and self.x_scaling_means is not None:
logger.debug("Saving input scaling information to %s_x_means.npy and %s_x_stds.npy", filename, filename)
np.save(filename + "_x_means.npy", self.x_scaling_means)
np.save(filename + "_x_stds.npy", self.x_scaling_stds)
np.save(filename + "_x_mins.npy", self.x_scaling_mins)
np.save(filename + "_x_maxs.npy", self.x_scaling_maxs)
# Save state dict
logger.debug("Saving state dictionary to %s_state_dict.pt", filename)
torch.save(self.model.state_dict(), filename + "_state_dict.pt")
# Save model
if save_model:
logger.debug("Saving model to %s_model.pt", filename)
torch.save(self.model, filename + "_model.pt")
# Export model to onnx
if export_model:
x = load_and_check(x)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dummy_input = torch.from_numpy(x[0].reshape(1, -1)).float().to(device)
torch.onnx.export(self.model, dummy_input,filename+".onnx", export_params=True, input_names = ['input'],output_names = ['r_hat', 's_hat'], verbose = True)
# Manipulate onnx model using 'onnxruntime' module directly
# Note: This is inefficient due to I/O reasons, however
# torch.onnx interface seemingly has no options for this
if export_model and os.path.isfile(filename+".onnx"):
####################################
## ONNXRUNTIME
## Example using Onnxruntime instead of Onxx
## Keeping only for prosperity for now
####################################
## Start the normal onnxruntime session using the model
## just saved
#ort_session = ort.InferenceSession(filename+".onnx")
## Model Meta data
#metaData = ort_session.get_modelmeta()
## Get the custom map
#CustomMap = metaData.custom_metadata_map
#print("Custom Meta-Data Map: {}".format(CustomMap))
## Define a new custom meta data map
#CustomMap_new = {"Var1" : 200.0,
# "Var2" : 5.0,
# "Var3" : 1000.0,
# "Var4" : 400.0,
# "Var5" : 6.0,
# }
#
## Load new custom map into mode
#metaData.custom_metadata_map = CustomMap_new
# Unable to save Onnx model from Onnxruntime Inference session it seems
# -> Makes sense given InferenceSession is designed to access and infer, not
# a data/model editor session.
#ort_session.SaveModelMetadata() # Believe that this does not work
####################################
####################################
####################################
## ONNX
####################################
# Define a new custom meta data map
#CustomMap_new = {"Var1" : 200.0,
# "Var2" : 5.0,
# "Var3" : 1000.0,
# "Var4" : 400.0,
# "Var5" : 6.0,
# }
# Load model
model = onnx.load(filename+".onnx")
# Get Meta Data
for index,(cust_key,cust_var) in enumerate(metaData.items()):
meta = model.metadata_props.add()
meta.key = cust_key
meta.value = str(cust_var)
# Check value
logger.info(" New Meta data: %s ",model.metadata_props[index])
# Save model
onnx.save(model, filename+"_new"+".onnx")
# Start the normal onnxruntime session using the model
# just saved to check that the model was saved with the correct
# metadata
ort_session = ort.InferenceSession(filename+"_new"+".onnx")
# Model Meta data
metaData = ort_session.get_modelmeta()
# Print Metadata
CustomMap = metaData.custom_metadata_map
logger.info(" Custom Meta-Data Map: %s",CustomMap)
# Need to close the ort session for comleteness (C-style)
####################################
###################################
# Tar model if training is done on GPU
if torch.cuda.is_available():
tar = tarfile.open("models_out.tar.gz", "w:gz")
for name in [filename+".onnx", filename + "_x_stds.npy", filename + "_x_means.npy", filename + "_x_mins.npy", filename + "_x_maxs.npy", filename + "_settings.json", filename + "_state_dict.pt"]:
tar.add(name)
tar.close()
import onnx
import sys
filename = "squeezenet.onnx"
if len(sys.argv) == 2:
filename = sys.argv[1]
print('Loading ONNX model %s' % filename)
# Load the ONNX model
model = onnx.load(filename)
# Check that the IR is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))
# Load pretrained ONNX model
# ---------------------------------------------
# The example super resolution model used here is exactly the same model in onnx tutorial
# http://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html
# we skip the pytorch model construction part, and download the saved onnx model
model_url = ''.join([
'https://gist.github.com/zhreshold/',
'bcda4716699ac97ea44f791c24310193/raw/',
'93672b029103648953c4e5ad3ac3aadf346a4cdc/', 'super_resolution_0.2.onnx'
])
#download(model_url, 'super_resolution.onnx', True)
download(model_url, 'super_resolution.onnx')
# now you have super_resolution.onnx on disk
onnx_model = onnx.load('super_resolution.onnx')
# we can load the graph as NNVM compatible model
sym, params = nnvm.frontend.from_onnx(onnx_model)
######################################################################
# Load a test image
# ---------------------------------------------
# A single cat dominates the examples!
from PIL import Image
img_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true'
download(img_url, 'cat.png')
img = Image.open('cat.png').resize((224, 224))
img_ycbcr = img.convert("YCbCr") # convert to YCbCr
img_y, img_cb, img_cr = img_ycbcr.split()
x = np.array(img_y)[np.newaxis, np.newaxis, :, :]
import os
import time
import sys
import onnx
from onnx import helper
import pdb
##
onnxFileName = str(sys.argv[1])
from onnx.shape_inference import infer_shapes
from onnx.optimizer import optimize
om = onnx.load(onnxFileName)
om.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1
om.graph.input[0].type.tensor_type.shape.dim[1].dim_value =3
om.graph.input[0].type.tensor_type.shape.dim[2].dim_value = 540
om.graph.input[0].type.tensor_type.shape.dim[3].dim_value = 960
v = om.graph.value_info
while len(v)>0:
v.pop()
#pdb.set_trace()
ots = om.graph.output
h = 17
for i in range(3)://this need do by yourself.
ots[i*4+0].type.tensor_type.shape.dim[2].dim_value=h
ots[i*4+1].type.tensor_type.shape.dim[2].dim_value=h
ots[i*4+2].type.tensor_type.shape.dim[2].dim_value=h
ots[i*4+3].type.tensor_type.shape.dim[2].dim_value=h
h *= 2
om = infer_shapes(om)
def load_model():
return onnx.load('../assets/emotion_ferplus/model.onnx')
# model URL, file name, and model type through the module, TVM will download
# the model and save it to disk. For the instance of an ONNX model, you can
# then load it into memory using the ONNX runtime.
#
# .. admonition:: Working with Other Model Formats
#
# TVM supports many popular model formats. A list can be found in the
# :ref:`Compile Deep Learning Models <tutorial-frontend>` section of the TVM
# Documentation.
model_url = ("https://github.com/onnx/models/raw/main/"
"vision/classification/resnet/model/"
"resnet50-v2-7.onnx")
model_path = download_testdata(model_url, "resnet50-v2-7.onnx", module="onnx")
onnx_model = onnx.load(model_path)
# Seed numpy's RNG to get consistent results
np.random.seed(0)
################################################################################
# Downloading, Preprocessing, and Loading the Test Image
# ------------------------------------------------------
#
# Each model is particular when it comes to expected tensor shapes, formats and
# data types. For this reason, most models require some pre and
# post-processing, to ensure the input is valid and to interpret the output.
# TVMC has adopted NumPy's ``.npz`` format for both input and output data.
#
# As input for this tutorial, we will use the image of a cat, but you can feel
# free to substitute this image for any of your choosing.
def onnx2tensorrt(onnx_file,
trt_file,
input_shape,
verify=False,
workspace_size=1):
"""Create tensorrt engine from onnx model.
Args:
onnx_file (str): Filename of the input ONNX model file.
trt_file (str): Filename of the output TensorRT engine file.
input_shape (list[int]): Input shape of the model.
eg [1, 3, 224, 224].
verify (bool, optional): Whether to verify the converted model.
Defaults to False.
workspace_size (int, optional): Maximium workspace of GPU.
Defaults to 1.
"""
import onnx
from mmcv.tensorrt import TRTWraper, onnx2trt, save_trt_engine
onnx_model = onnx.load(onnx_file)
# create trt engine and wraper
opt_shape_dict = {'input': [input_shape, input_shape, input_shape]}
max_workspace_size = get_GiB(workspace_size)
trt_engine = onnx2trt(
onnx_model,
opt_shape_dict,
fp16_mode=False,
max_workspace_size=max_workspace_size)
save_dir, _ = osp.split(trt_file)
if save_dir:
os.makedirs(save_dir, exist_ok=True)
save_trt_engine(trt_engine, trt_file)
print(f'Successfully created TensorRT engine: {trt_file}')
if verify:
import torch
import onnxruntime as ort
input_img = torch.randn(*input_shape)
input_img_cpu = input_img.detach().cpu().numpy()
input_img_cuda = input_img.cuda()
# Get results from ONNXRuntime
session_options = ort.SessionOptions()
sess = ort.InferenceSession(onnx_file, session_options)
# get input and output names
input_names = [_.name for _ in sess.get_inputs()]
output_names = [_.name for _ in sess.get_outputs()]
onnx_outputs = sess.run(None, {
input_names[0]: input_img_cpu,
})
# Get results from TensorRT
trt_model = TRTWraper(trt_file, input_names, output_names)
with torch.no_grad():
trt_outputs = trt_model({input_names[0]: input_img_cuda})
trt_outputs = [
trt_outputs[_].detach().cpu().numpy() for _ in output_names
]
# Compare results
np.testing.assert_allclose(
onnx_outputs[0], trt_outputs[0], rtol=1e-05, atol=1e-05)
print('The numerical values are the same ' +
'between ONNXRuntime and TensorRT')
import onnx
import sys
# Load the ONNX model
#model = onnx.load("alexnet.proto")
modelName = "torch_linear.proto"
if len(sys.argv) > 1:
modelName = sys.argv[1]
model = onnx.load(modelName)
onnx.checker.check_model(model)
print(onnx.helper.printable_graph(model.graph))
import sys
import timeit
import numpy as np
import onnx
import ngraph as ng
from ngraph_onnx.onnx_importer.importer import import_onnx_model
model = onnx.load(sys.argv[1])
ng_func = import_onnx_model(model)
#print(ng_model)
picture = np.ones([1, 3, 224, 224], dtype=np.float32)
runtime = ng.runtime(backend_name='CPU')
#runtime = ng.runtime(backend_name='GPU')
resnet = runtime.computation(ng_func)
#print(resnet)
def run():
resnet(picture)
n = 100
print(timeit.timeit('run()', globals=globals(), number=n) / n * 1000, 'msec')
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ******************************************************************************
import onnx
onnx_protobuf = onnx.load('/path/to/model/cntk_ResNet20_CIFAR10/model.onnx')
# Convert a serialized ONNX model to an ngraph model
from ngraph_onnx.onnx_importer.importer import import_onnx_model
ng_model = import_onnx_model(onnx_protobuf)[0]
# Using an ngraph runtime (CPU backend), create a callable computation
import ngraph as ng
runtime = ng.runtime(backend_name='CPU')
resnet = runtime.computation(ng_model['output'], *ng_model['inputs'])
# Load or create an image
import numpy as np
picture = np.ones([1, 3, 32, 32])
def pytorch2onnx(model,
mm_inputs,
opset_version=11,
show=False,
output_file='tmp.onnx',
verify=False,
dynamic_export=False):
"""Export Pytorch model to ONNX model and verify the outputs are same
between Pytorch and ONNX.
Args:
model (nn.Module): Pytorch model we want to export.
mm_inputs (dict): Contain the input tensors and img_metas information.
opset_version (int): The onnx op version. Default: 11.
show (bool): Whether print the computation graph. Default: False.
output_file (string): The path to where we store the output ONNX model.
Default: `tmp.onnx`.
verify (bool): Whether compare the outputs between Pytorch and ONNX.
Default: False.
dynamic_export (bool): Whether to export ONNX with dynamic axis.
Default: False.
"""
model.cpu().eval()
test_mode = model.test_cfg.mode
if isinstance(model.decode_head, nn.ModuleList):
num_classes = model.decode_head[-1].num_classes
else:
num_classes = model.decode_head.num_classes
imgs = mm_inputs.pop('imgs')
img_metas = mm_inputs.pop('img_metas')
img_list = [img[None, :] for img in imgs]
img_meta_list = [[img_meta] for img_meta in img_metas]
# update img_meta
img_list, img_meta_list = _update_input_img(img_list, img_meta_list)
# replace original forward function
origin_forward = model.forward
model.forward = partial(model.forward,
img_metas=img_meta_list,
return_loss=False,
rescale=True)
dynamic_axes = None
if dynamic_export:
if test_mode == 'slide':
dynamic_axes = {'input': {0: 'batch'}, 'output': {1: 'batch'}}
else:
dynamic_axes = {
'input': {
0: 'batch',
2: 'height',
3: 'width'
},
'output': {
1: 'batch',
2: 'height',
3: 'width'
}
}
register_extra_symbolics(opset_version)
with torch.no_grad():
torch.onnx.export(model, (img_list, ),
output_file,
input_names=['input'],
output_names=['output'],
export_params=True,
keep_initializers_as_inputs=False,
verbose=show,
opset_version=opset_version,
dynamic_axes=dynamic_axes)
print(f'Successfully exported ONNX model: {output_file}')
model.forward = origin_forward
if verify:
# check by onnx
import onnx
onnx_model = onnx.load(output_file)
onnx.checker.check_model(onnx_model)
if dynamic_export and test_mode == 'whole':
# scale image for dynamic shape test
img_list = [resize(_, scale_factor=1.5) for _ in img_list]
# concate flip image for batch test
flip_img_list = [_.flip(-1) for _ in img_list]
img_list = [
torch.cat((ori_img, flip_img), 0)
for ori_img, flip_img in zip(img_list, flip_img_list)
]
# update img_meta
img_list, img_meta_list = _update_input_img(
img_list, img_meta_list, test_mode == 'whole')
# check the numerical value
# get pytorch output
with torch.no_grad():
pytorch_result = model(img_list, img_meta_list, return_loss=False)
pytorch_result = np.stack(pytorch_result, 0)
# get onnx output
input_all = [node.name for node in onnx_model.graph.input]
input_initializer = [
node.name for node in onnx_model.graph.initializer
]
net_feed_input = list(set(input_all) - set(input_initializer))
assert (len(net_feed_input) == 1)
sess = rt.InferenceSession(output_file)
onnx_result = sess.run(
None, {net_feed_input[0]: img_list[0].detach().numpy()})[0][0]
# show segmentation results
if show:
import cv2
import os.path as osp
img = img_meta_list[0][0]['filename']
if not osp.exists(img):
img = imgs[0][:3, ...].permute(1, 2, 0) * 255
img = img.detach().numpy().astype(np.uint8)
ori_shape = img.shape[:2]
else:
ori_shape = LoadImage()({'img': img})['ori_shape']
# resize onnx_result to ori_shape
onnx_result_ = cv2.resize(onnx_result[0].astype(np.uint8),
(ori_shape[1], ori_shape[0]))
show_result_pyplot(model,
img, (onnx_result_, ),
palette=model.PALETTE,
block=False,
title='ONNXRuntime',
opacity=0.5)
# resize pytorch_result to ori_shape
pytorch_result_ = cv2.resize(pytorch_result[0].astype(np.uint8),
(ori_shape[1], ori_shape[0]))
show_result_pyplot(model,
img, (pytorch_result_, ),
title='PyTorch',
palette=model.PALETTE,
opacity=0.5)
# compare results
np.testing.assert_allclose(
pytorch_result.astype(np.float32) / num_classes,
onnx_result.astype(np.float32) / num_classes,
rtol=1e-5,
atol=1e-5,
err_msg='The outputs are different between Pytorch and ONNX')
print('The outputs are same between Pytorch and ONNX')
type=str,
help='Path to text file containing the names of all possible categories.')
parser.add_argument('--image_path', type=str, help='Path to image file.')
parser.add_argument('--image_list', type=str, help='Path to image list.')
args = parser.parse_args()
IMAGE_FILENAME = args.image_path # '/home/meerkat/git/SpeciesClassification/PyTorchClassification/elephant.jpg'
MODEL_FILENAME = args.frozen_graph # '/ai4edevfs/models/iNat/iNat_all_extended/demosite-model-ensemble-resnext-inceptionV4-560-81.0/iNat_all_extended_ensemble_resnext_inceptionV4_560_81.9_model.onnx'
CLASSLIST_FILENAME = args.classlist # '/ai4edevfs/models/iNat/iNat_all_extended/demosite-model-ensemble-resnext-inceptionV4-560-81.0/iNat_all_extended_ensemble_resnext_inceptionV4_560_81.9_classes.txt'
# Target mean / std; should match the values used at training time
MODEL_IMAGE_SIZE = 560
MODEL_RESIZE_SIZE = 640
OVERSIZE_FACTOR = 1.3
model = onnx.load(MODEL_FILENAME)
if args.image_list:
with open(args.image_list, 'rt') as fi:
all_images = fi.read().splitlines()
else:
all_images = [IMAGE_FILENAME]
for IMAGE_FILENAME in all_images:
if not os.path.isfile(IMAGE_FILENAME):
print("Did not find " + IMAGE_FILENAME)
continue
#%% Load and prepare image (using cv2)
def detect(save_img=False):
imgsz = (
320, 192
) if ONNX_EXPORT else opt.img_size # (320, 192) or (416, 256) or (608, 352) for (height, width)
out, source, weights, half, view_img, save_txt = opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = torch_utils.select_device(
device='cpu' if ONNX_EXPORT else opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Initialize model
model = Darknet(opt.cfg, imgsz)
# Load weights
attempt_download(weights)
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(
torch.load(weights, map_location=device)['model'])
else: # darknet format
load_darknet_weights(model, weights)
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']) # load weights
modelc.to(device).eval()
# Eval mode
model.to(device).eval()
# Fuse Conv2d + BatchNorm2d layers
# model.fuse()
# Export mode
if ONNX_EXPORT:
model.fuse()
img = torch.zeros((1, 3) + imgsz) # (1, 3, 320, 192)
f = opt.weights.replace(opt.weights.split('.')[-1],
'onnx') # *.onnx filename
torch.onnx.export(model,
img,
f,
verbose=False,
opset_version=11,
input_names=['images'],
output_names=['classes', 'boxes'])
# Validate exported model
import onnx
model = onnx.load(f) # Load the ONNX model
onnx.checker.check_model(model) # Check that the IR is well formed
print(onnx.helper.printable_graph(
model.graph)) # Print a human readable representation of the graph
return
# Half precision
half = half and device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = load_classes(opt.names)
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img.float()
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt.augment)[0]
t2 = torch_utils.time_synchronized()
# to float
if half:
pred = pred.float()
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
multi_label=False,
classes=opt.classes,
agnostic=opt.agnostic_nms)
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections for image i
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from imgsz to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(save_path[:save_path.rfind('.')] + '.txt',
'a') as file:
file.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)])
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def pytorch2onnx(model: nn.Module,
model_type: str,
img_path: str,
verbose: bool = False,
show: bool = False,
opset_version: int = 11,
output_file: str = 'tmp.onnx',
verify: bool = False,
dynamic_export: bool = False,
device_id: int = 0):
"""Export Pytorch model to ONNX model and verify the outputs are same
between Pytorch and ONNX.
Args:
model (nn.Module): Pytorch model we want to export.
model_type (str): Model type, detection or recognition model.
img_path (str): We need to use this input to execute the model.
opset_version (int): The onnx op version. Default: 11.
verbose (bool): Whether print the computation graph. Default: False.
show (bool): Whether visialize final results. Default: False.
output_file (string): The path to where we store the output ONNX model.
Default: `tmp.onnx`.
verify (bool): Whether compare the outputs between Pytorch and ONNX.
Default: False.
dynamic_export (bool): Whether apply dynamic export.
Default: False.
device_id (id): Device id to place model and data.
Default: 0
"""
device = torch.device(type='cuda', index=device_id)
model.to(device).eval()
_convert_batchnorm(model)
# prepare inputs
mm_inputs = _prepare_data(cfg=model.cfg, imgs=img_path)
imgs = mm_inputs.pop('img')
img_metas = mm_inputs.pop('img_metas')
if isinstance(imgs, list):
imgs = imgs[0]
img_list = [img[None, :].to(device) for img in imgs]
# update img_meta
img_list, img_metas = _update_input_img(img_list, img_metas)
origin_forward = model.forward
if (model_type == 'det'):
model.forward = partial(model.simple_test,
img_metas=img_metas,
rescale=True)
else:
model.forward = partial(model.forward,
img_metas=img_metas,
return_loss=False,
rescale=True)
# pytorch has some bug in pytorch1.3, we have to fix it
# by replacing these existing op
register_extra_symbolics(opset_version)
dynamic_axes = None
if dynamic_export and model_type == 'det':
dynamic_axes = {
'input': {
0: 'batch',
2: 'height',
3: 'width'
},
'output': {
0: 'batch',
2: 'height',
3: 'width'
}
}
elif dynamic_export and model_type == 'recog':
dynamic_axes = {
'input': {
0: 'batch',
3: 'width'
},
'output': {
0: 'batch',
3: 'width'
}
}
with torch.no_grad():
torch.onnx.export(model, (img_list[0], ),
output_file,
input_names=['input'],
output_names=['output'],
export_params=True,
keep_initializers_as_inputs=False,
verbose=verbose,
opset_version=opset_version,
dynamic_axes=dynamic_axes)
print(f'Successfully exported ONNX model: {output_file}')
if verify:
# check by onnx
import onnx
onnx_model = onnx.load(output_file)
onnx.checker.check_model(onnx_model)
scale_factor = (0.5, 0.5) if model_type == 'det' else (1, 0.5)
if dynamic_export:
# scale image for dynamic shape test
img_list = [
nn.functional.interpolate(_, scale_factor=scale_factor)
for _ in img_list
]
# update img_meta
img_list, img_metas = _update_input_img(img_list, img_metas)
# check the numerical value
# get pytorch output
with torch.no_grad():
model.forward = origin_forward
pytorch_out = model.simple_test(img_list[0],
img_metas[0],
rescale=True)
# get onnx output
if model_type == 'det':
onnx_model = ONNXRuntimeDetector(output_file, model.cfg, device_id)
else:
onnx_model = ONNXRuntimeRecognizer(output_file, model.cfg,
device_id)
onnx_out = onnx_model.simple_test(img_list[0],
img_metas[0],
rescale=True)
# compare results
same_diff = 'same'
if model_type == 'recog':
for onnx_result, pytorch_result in zip(onnx_out, pytorch_out):
if onnx_result['text'] != pytorch_result[
'text'] or not np.allclose(
np.array(onnx_result['score']),
np.array(pytorch_result['score']),
rtol=1e-4,
atol=1e-4):
same_diff = 'different'
break
else:
for onnx_result, pytorch_result in zip(
onnx_out[0]['boundary_result'],
pytorch_out[0]['boundary_result']):
if not np.allclose(np.array(onnx_result),
np.array(pytorch_result),
rtol=1e-4,
atol=1e-4):
same_diff = 'different'
break
print('The outputs are {} between Pytorch and ONNX'.format(same_diff))
if show:
onnx_img = onnx_model.show_result(img_path,
onnx_out[0],
out_file='onnx.jpg',
show=False)
pytorch_img = model.show_result(img_path,
pytorch_out[0],
out_file='pytorch.jpg',
show=False)
if onnx_img is None:
onnx_img = cv2.imread(img_path)
if pytorch_img is None:
pytorch_img = cv2.imread(img_path)
cv2.imshow('PyTorch', pytorch_img)
cv2.imshow('ONNXRuntime', onnx_img)
cv2.waitKey()
return
import urllib
urllib.urlretrieve(url, path)
######################################################################
# Load pretrained ONNX model
# ---------------------------------------------
# The example super resolution model used here is exactly the same model in onnx tutorial
# http://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html
# we skip the pytorch model construction part, and download the saved onnx model
model_url = ''.join(['https://gist.github.com/zhreshold/',
'bcda4716699ac97ea44f791c24310193/raw/',
'93672b029103648953c4e5ad3ac3aadf346a4cdc/',
'super_resolution_0.2.onnx'])
download(model_url, 'super_resolution.onnx', True)
# now you have super_resolution.onnx on disk
onnx_graph = onnx.load('super_resolution.onnx')
# we can load the graph as NNVM compatible model
sym, params = nnvm.frontend.from_onnx(onnx_graph)
######################################################################
# Load a test image
# ---------------------------------------------
# A single cat dominates the examples!
from PIL import Image
img_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true'
download(img_url, 'cat.png')
img = Image.open('cat.png').resize((224, 224))
img_ycbcr = img.convert("YCbCr") # convert to YCbCr
img_y, img_cb, img_cr = img_ycbcr.split()
x = np.array(img_y)[np.newaxis, np.newaxis, :, :]
import onnx
from onnxsim import simplify
input_path= '/data/ai/yolov5/pet2kyolov5n.onnx'
output_path='/data/ai/yolov5/pet2kyolov5n_simpy.onnx'
onnx_model = onnx.load(input_path) # load onnx model
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, output_path)
print('finished exporting onnx')
import numpy as np
import onnx
import onnxruntime
from tvm import relay
import tvm
onnx_model = onnx.load('05_dense_b.onnx')
input_name1 = 'X'
input_shape1 = (1, 1)
shape_dict = {input_name1: input_shape1}
mod, params = relay.frontend.from_onnx(model=onnx_model, shape=shape_dict)
# Compilation
opt_level = 3
target = 'llvm'
with relay.build_config(opt_level=opt_level):
graph, lib, params = relay.build_module.build(mod, target, params=params)
#printing some LLVM code
out_file = open("05_dense_b.ll", "w")
out_file.write(lib.get_source())
out_file.close()
bias=bias,
dilation=dilation,
groups=groups)
model.eval()
input_np = np.random.uniform(
0, 1, (1, 3, dimension, 224, 224)) # NCDHW
input_var = Variable(
torch.FloatTensor(input_np))
torch.onnx.export(model,
input_var,
"_tmpnet.onnx",
verbose=True,
input_names=['test_in'],
output_names=['test_out'])
onnx_model = onnx.load('_tmpnet.onnx')
k_model = onnx_to_keras(
onnx_model, ['test_in'])
error = check_torch_keras_error(
model, k_model, input_np)
print('Error:', error)
if max_error < error:
max_error = error
print('Max error: {0}'.format(max_error))
input_names=['images'],
dynamic_axes={
'images': {
0: 'batch',
2: 'height',
3: 'width'
}, # size(1,3,640,640)
'output': {
0: 'batch',
2: 'y',
3: 'x'
}
} if opt.dynamic else None)
# Checks
model_onnx = onnx.load(f) # load onnx model
onnx.checker.check_model(model_onnx) # check onnx model
# print(onnx.helper.printable_graph(model_onnx.graph)) # print
# Simplify
if opt.simplify:
try:
check_requirements(['onnx-simplifier'])
import onnxsim
print(
f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...'
)
model_onnx, check = onnxsim.simplify(
model_onnx,
dynamic_input_shape=opt.dynamic,
# can be expressed using Gemm alternatively
node = oh.make_node('MatMul', ["pdf","cdfmat"], ["cdf"])
# ArgMax(Clip(cdf + (1 - rand), max=1))
# node_rand_sub = oh.make_node('Sub', ["one","r"], ["rr"])
add_node = oh.make_node('Add', ["cdf", "r"], ["temp"], broadcast=1)
clip_node = oh.make_node('Clip', ["temp"], ["cdf_clipped"], min=0.0, max=1.0)
topk_node = oh.make_node('ArgMax', ["cdf_clipped"], ["chosen_action"], axis=1, keepdims=0)
graph = oh.make_graph([one, one_int, pdf_temp, pdf_temp_2,
max_score, normalized_scores, # top_action_1, # top_action_2,
# one_hot_top_action_int, one_hot_top_action_float,
one_hot_top_action,
exploit_prob, exploit_top_action,
pdf, # node_rand_sub,
node, node_rand, add_node, clip_node, topk_node],
'compute_graph', input_tensors, output_tensors, initializer_tensors)
model = oh.make_model(graph, producer_name='explore')
f = open("model1.onnx", "wb")
f.write(model.SerializeToString())
f.close()
tf_model = onnx.load('model1.onnx')
tf_rep = prepare(tf_model)
sample = tf_rep.run(np.asarray([[.3,.3,.4]]))
print(sample)
import onnx
# Load the ONNX model
model = onnx.load('./output/hopenet.pb')
# Check that the IR is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))
if len(sys.argv) != 4:
raise RuntimeError('Usage: %s input_dir output_dir ratio' % sys.argv[0])
input_dir = sys.argv[1]
output_dir = sys.argv[2]
ratio = float(sys.argv[3])
os.makedirs(output_dir, exist_ok=True)
shutil.rmtree(os.path.join(output_dir, 'test_data_set_0'), ignore_errors=True)
shutil.copytree(os.path.join(input_dir, 'test_data_set_0'),
os.path.join(output_dir, 'test_data_set_0'))
model = onnx.load(os.path.join(input_dir, 'model.onnx'))
initializer_names = {i.name for i in model.graph.initializer}
inputs = []
for input in model.graph.input:
if input.name not in initializer_names:
inputs.append(input)
assert len(inputs) == 1
input = inputs[0]
old_shape = []
new_shape = []
for i in range(len(input.type.tensor_type.shape.dim)):
nd = input.type.tensor_type.shape.dim[i].dim_value
old_shape.append(nd)
def simplify(model: Union[str, onnx.ModelProto],
inputs: Sequence[Dict[str, np.ndarray]] = None,
output_file: str = None,
perform_optimization: bool = True,
skip_fuse_bn: bool = False,
skip_shape_inference: bool = True,
input_shapes: Dict[str, Sequence[int]] = None,
skipped_optimizers: Sequence[str] = None) -> onnx.ModelProto:
"""Simplify and optimize an onnx model.
For models from detection and segmentation, it is strongly suggested to
input multiple input images for verification.
Arguments:
model (str or onnx.ModelProto): path of model or loaded model object.
inputs (optional, Sequence[Dict[str, np.ndarray]]): inputs of model.
output_file (optional, str): output file to save simplified model.
perform_optimization (optional, bool): whether to perform optimization.
skip_fuse_bn (optional, bool): whether to skip fusing bn layer.
skip_shape_inference (optional, bool): whether to skip shape inference.
input_shapes (optional, Dict[str, Sequence[int]]):
the shapes of model inputs.
skipped_optimizers (optional, Sequence[str]):
the names of optimizer to be skipped.
Returns:
onnx.ModelProto: simplified and optimized onnx model.
Example:
>>> import onnx
>>> import numpy as np
>>>
>>> from mmcv.onnx import simplify
>>>
>>> dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
>>> input = {'input':dummy_input}
>>> input_file = 'sample.onnx'
>>> output_file = 'slim.onnx'
>>> model = simplify(input_file, [input], output_file)
"""
if input_shapes is None:
input_shapes = {}
if isinstance(model, str):
model = onnx.load(model)
# rename op with numeric name for issue
# https://github.com/onnx/onnx/issues/2613
model = add_suffix2name(model)
onnx.checker.check_model(model)
model_ori = copy.deepcopy(model)
numel_node_ori = len(model_ori.graph.node)
if not skip_shape_inference:
model = onnx.shape_inference.infer_shapes(model)
input_shapes = check_and_update_input_shapes(model, input_shapes)
if perform_optimization:
model = optimize(model, skip_fuse_bn, skipped_optimizers)
const_nodes = get_constant_nodes(model)
feed_inputs = None if inputs is None else inputs[0]
res = forward_for_node_outputs(model,
const_nodes,
input_shapes=input_shapes,
inputs=feed_inputs)
const_nodes = clean_constant_nodes(const_nodes, res)
model = eliminate_const_nodes(model, const_nodes, res)
onnx.checker.check_model(model)
if perform_optimization:
model = optimize(model, skip_fuse_bn, skipped_optimizers)
check_ok = check(model_ori,
model,
input_shapes=input_shapes,
inputs=inputs)
assert check_ok, 'Check failed for the simplified model!'
numel_node_slim = len(model.graph.node)
print('Number of nodes: {numel_node_ori} -> {numel_node_slim}')
if output_file is not None:
save_dir, _ = os.path.split(output_file)
if save_dir:
os.makedirs(save_dir, exist_ok=True)
onnx.save(model, output_file)
return model
# the sequence. the second is just the most recent hidden state
# (compare the last slice of "out" with "hidden" below, they are the same)
# The reason for this is that:
# "out" will give you access to all hidden states in the sequence
# "hidden" will allow you to continue the sequence and backpropagate,
# by passing it as an argument to the lstm at a later time
# Add the extra 2nd dimension
out, hidden = lstm(TEST_INPUTS, (TEST_INITIAL_H, TEST_INITIAL_C))
#out, hidden = lstm(TEST_INPUTS, (TEST_INITIAL_H_2, TEST_INITIAL_C_2))
print(out)
print(hidden)
torch.onnx.export(lstm, (TEST_INPUTS, (TEST_INITIAL_H, TEST_INITIAL_C)),
"lstm.onnx",
verbose=False)
else:
# read the model from lstm.onnx with onnx-tensorflow
model = onnx.load("lstm.onnx")
tf_rep = prepare(model)
import tensorflow as tf
print(
tf_rep.run({
"0": TEST_INPUTS,
"1": TEST_INITIAL_H,
"2": TEST_INITIAL_C
}))
def convertAndSave(path, outputPath):
o = onnx.load(path)
onnxInputLayer, = set(map(lambda x: x.name, o.graph.input)) - set(map(lambda x: x.name, o.graph.initializer))
# Parse layers
layers = lbann.onnx.o2l.onnxToLbannLayers(
o,
[dataLayerName, labelLayerName],
{dataLayerName: onnxInputLayer},
)
# Add a softmax layer
outputLayerName = layers[-1].name
probLayerName = outputLayerName
# layers.append(lbann_pb2.Layer(name=probLayerName,
# parents=lbann.onnx.util.list2LbannList([outputLayerName]),
# data_layout="data_parallel",
# softmax=lbann_pb2.Softmax()))
# Add metric layers
for name, dic in [(crossEntropyLayerName, {"cross_entropy": lbann_pb2.CrossEntropy()}),
(top1AccuracyLayerName, {"categorical_accuracy": lbann_pb2.CategoricalAccuracy()}),
(top5AccuracyLayerName, {"top_k_categorical_accuracy": lbann_pb2.TopKCategoricalAccuracy(k=5)})]:
layers.append(lbann_pb2.Layer(name=name,
parents=lbann.onnx.util.list2LbannList([probLayerName, labelLayerName]),
data_layout="data_parallel",
**dic))
# Define an objective function
objective = lbann_pb2.ObjectiveFunction(
layer_term = [lbann_pb2.LayerTerm(layer=crossEntropyLayerName)],
l2_weight_regularization = [lbann_pb2.L2WeightRegularization(scale_factor=1e-4)]
)
# Add metrics
metrics = []
for name, layer, unit in [("categorical accuracy", top1AccuracyLayerName, "%"),
("top-5 categorical accuracy", top5AccuracyLayerName, "%")]:
metrics.append(lbann_pb2.Metric(layer_metric=lbann_pb2.LayerMetric(name=name,
layer=layer,
unit=unit)))
# Add callbacks
callbacks = []
for dic in [{"print": lbann_pb2.CallbackPrint()},
{"timer": lbann_pb2.CallbackTimer()},
{"imcomm": lbann_pb2.CallbackImComm(intermodel_comm_method="normal", all_optimizers=True)}]:
callbacks.append(lbann_pb2.Callback(**dic))
model = lbann_pb2.Model(
data_layout = "data_parallel",
mini_batch_size = 256,
block_size = 256,
num_epochs = 10,
num_parallel_readers = 0,
procs_per_model = 0,
objective_function = objective,
metric = metrics,
callback = callbacks,
layer = layers
)
with open(outputPath, "w") as f:
f.write(txtf.MessageToString(lbann_pb2.LbannPB(model=model)))
def detect(save_txt=False, save_img=False):
# (320, 192) or (416, 256) or (608, 352) for (height, width)
# img_size = (320, 192) if ONNX_EXPORT else opt.img_size
img_size = (352, 608)
out, source, weights, half, view_img = opt.output, opt.source, opt.weights, opt.half, opt.view_img
webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
# Initialize
device = torch_utils.select_device(device='cpu' if ONNX_EXPORT else opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Initialize model
model = Darknet(opt.cfg, img_size)
# Load weights
attempt_download(weights)
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(torch.load(weights, map_location=device)['model'])
else: # darknet format
_ = load_darknet_weights(model, weights)
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101', n=2) # initialize
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights
modelc.to(device).eval()
# Fuse Conv2d + BatchNorm2d layers
# model.fuse()
# Eval mode
model.to(device).eval()
# Export mode
if ONNX_EXPORT:
img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192)
torch.onnx.export(model, img, 'weights/export.onnx', verbose=True, opset_version=11)
# Validate exported model
import onnx
model = onnx.load('weights/export.onnx') # Load the ONNX model
onnx.checker.check_model(model) # Check that the IR is well formed
print(onnx.helper.printable_graph(model.graph)) # Print a human readable representation of the graph
return
# Half precision
half = half and device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=img_size, half=half)
else:
save_img = True
dataset = LoadImages(source, img_size=img_size, half=half)
# Get names and colors
names = load_classes(opt.names)
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
# Run inference
timeStr = None
t0 = time.time()
for path, img, im0s, vid_cap in dataset:
t = time.time()
match = re.match(".*T([^\\.]*)", path)
if match is not None:
groups = match.groups()
if len(groups) > 0:
timeStr = groups[0].replace("-", ":")
# Get detections
img = torch.from_numpy(img).to(device)
if img.ndimension() == 3:
img = img.unsqueeze(0)
pred = model(img)[0]
if opt.half:
pred = pred.float()
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.nms_thres)
# Apply
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i]
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
s += '%gx%g ' % img.shape[2:] # print string
people = {}
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
with open(save_path + '.txt', 'a') as file:
file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf))
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])
if names[int(cls)] == "person":
for threshold in range(90, 0, -10):
if conf.item() * 100 > threshold:
people[str(threshold)] = people.get(str(threshold), 0) + 1
line = f"{timeStr},"
for threshold in range(90, 0, -10):
line += f"{people.get(str(threshold), 0)},"
file.write(f"{line}\n")
file.flush()
print('%sDone. (%.3fs)' % (s, time.time() - t))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + out + ' ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def convert_onnx_model(model_path, onnx_file, model_name):
"""
Util to convert onnx model to MXNet model
:param model_name:
:param model_path:
:param onnx_file:
:return:
"""
try:
import mxnet as mx
from mxnet.contrib import onnx as onnx_mxnet
except ImportError:
raise ModelArchiverError(
"MXNet package is not installed. Run command: pip install mxnet to install it."
)
try:
import onnx
except ImportError:
raise ModelArchiverError(
"Onnx package is not installed. Run command: pip install onnx to install it."
)
symbol_file = '%s-symbol.json' % model_name
params_file = '%s-0000.params' % model_name
signature_file = 'signature.json'
# Find input symbol name and shape
try:
model_proto = onnx.load(os.path.join(model_path, onnx_file))
except:
logging.error(
"Failed to load the %s model. Verify if the model file is valid",
onnx_file)
raise
graph = model_proto.graph
_params = set()
for tensor_vals in graph.initializer:
_params.add(tensor_vals.name)
input_data = []
for graph_input in graph.input:
shape = []
if graph_input.name not in _params:
for val in graph_input.type.tensor_type.shape.dim:
shape.append(val.dim_value)
input_data.append((graph_input.name, tuple(shape)))
try:
sym, arg_params, aux_params = onnx_mxnet.import_model(
os.path.join(model_path, onnx_file))
# UNION of argument and auxillary parameters
params = dict(arg_params, **aux_params)
except:
logging.error(
"Failed to import %s file to onnx. Verify if the model file is valid",
onnx_file)
raise
try:
# rewrite input data_name correctly
with open(os.path.join(model_path, signature_file), 'r') as f:
data = json.loads(f.read())
data['inputs'][0]['data_name'] = input_data[0][0]
data['inputs'][0]['data_shape'] = [
int(i) for i in input_data[0][1]
]
with open(os.path.join(model_path, signature_file), 'w') as f:
f.write(json.dumps(data, indent=2))
with open(os.path.join(model_path, symbol_file), 'w') as f:
f.write(sym.tojson())
except:
logging.error(
"Failed to write the signature or symbol files for %s model",
onnx_file)
raise
save_dict = {('arg:%s' % k): v.as_in_context(mx.cpu())
for k, v in params.items()}
mx.nd.save(os.path.join(model_path, params_file), save_dict)
return symbol_file, params_file
import sys
import onnx
import onnx.shape_inference
if len(sys.argv) != 2:
raise RuntimeError('Usage: %s input.onnx' % sys.argv[0])
model = onnx.load(sys.argv[1], 'rb')
model = onnx.shape_inference.infer_shapes(model)
value_infos = {}
for vi in model.graph.input:
value_infos[vi.name] = vi
for vi in model.graph.output:
value_infos[vi.name] = vi
for vi in model.graph.value_info:
value_infos[vi.name] = vi
total_kflops = 0
for node in model.graph.node:
if node.op_type != 'Conv':
continue
if (node.input[0] in value_infos and
node.input[1] in value_infos and
node.output[0] in value_infos):
in_shape = value_infos[node.input[0]].type.tensor_type.shape
w_shape = value_infos[node.input[1]].type.tensor_type.shape
out_shape = value_infos[node.output[0]].type.tensor_type.shape
kvs = [
def create_test_dir(model_path,
root_path,
test_name,
name_input_map=None,
symbolic_dim_values_map=None,
name_output_map=None):
"""
Create a test directory that can be used with onnx_test_runner or onnxruntime_perf_test.
Generates random input data for any missing inputs.
Saves output from running the model if name_output_map is not provided.
:param model_path: Path to the onnx model file to use.
:param root_path: Root path to create the test directory in.
:param test_name: Name for test. Will be added to the root_path to create the test directory name.
:param name_input_map: Map of input names to numpy ndarray data for each input.
:param symbolic_dim_values_map: Map of symbolic dimension names to values to use for the input data if creating
using random data.
:param name_output_map: Optional map of output names to numpy ndarray expected output data.
If not provided, the model will be run with the input to generate output data to save.
:return: None
"""
model_path = os.path.abspath(model_path)
root_path = os.path.abspath(root_path)
test_dir = os.path.join(root_path, test_name)
if not os.path.exists(test_dir):
os.makedirs(test_dir)
# add to existing test data sets if present
test_num = 0
while True:
test_data_dir = os.path.join(test_dir,
"test_data_set_" + str(test_num))
if not os.path.exists(test_data_dir):
os.mkdir(test_data_dir)
break
test_num += 1
model_filename = os.path.split(model_path)[-1]
test_model_filename = os.path.join(test_dir, model_filename)
shutil.copy(model_path, test_model_filename)
model = onnx.load(model_path)
model_inputs = model.graph.input
model_outputs = model.graph.output
def save_data(prefix, name_data_map, model_info):
idx = 0
for name, data in name_data_map.items():
if isinstance(data, dict):
# ignore. map<T1, T2> from traditional ML ops
pass
elif isinstance(data, list):
# ignore. vector<map<T1,T2>> from traditional ML ops. e.g. ZipMap output
pass
else:
np_type = _get_numpy_type(model_info, name)
tensor = numpy_helper.from_array(data.astype(np_type), name)
filename = os.path.join(test_data_dir,
"{}_{}.pb".format(prefix, idx))
with open(filename, "wb") as f:
f.write(tensor.SerializeToString())
idx += 1
if not name_input_map:
name_input_map = {}
if not symbolic_dim_values_map:
symbolic_dim_values_map = {}
initializer_set = set()
for initializer in onnx.load(model_path).graph.initializer:
initializer_set.add(initializer.name)
_create_missing_input_data(model_inputs, name_input_map,
symbolic_dim_values_map, initializer_set)
save_data("input", name_input_map, model_inputs)
# save expected output data if provided. run model to create if not.
if not name_output_map:
output_names = [o.name for o in model_outputs]
sess = ort.InferenceSession(test_model_filename)
outputs = sess.run(output_names, name_input_map)
name_output_map = {}
for name, data in zip(output_names, outputs):
name_output_map[name] = data
save_data("output", name_output_map, model_outputs)
import tvm.relay as relay
from tvm.contrib.download import download_testdata
######################################################################
# Load pretrained ONNX model
# ---------------------------------------------
# The example super resolution model used here is exactly the same model in onnx tutorial
# http://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html
# we skip the pytorch model construction part, and download the saved onnx model
model_url = ''.join(['https://gist.github.com/zhreshold/',
'bcda4716699ac97ea44f791c24310193/raw/',
'93672b029103648953c4e5ad3ac3aadf346a4cdc/',
'super_resolution_0.2.onnx'])
model_path = download_testdata(model_url, 'super_resolution.onnx', module='onnx')
# now you have super_resolution.onnx on disk
onnx_model = onnx.load(model_path)
######################################################################
# Load a test image
# ---------------------------------------------
# A single cat dominates the examples!
from PIL import Image
img_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true'
img_path = download_testdata(img_url, 'cat.png', module='data')
img = Image.open(img_path).resize((224, 224))
img_ycbcr = img.convert("YCbCr") # convert to YCbCr
img_y, img_cb, img_cr = img_ycbcr.split()
x = np.array(img_y)[np.newaxis, np.newaxis, :, :]
######################################################################
# Compile the model with relay
import onnxruntime
from onnxsim import simplify
import argparse
import os
from pprint import pprint
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, required=True)
parser.add_argument("--height", type=int, required=True)
parser.add_argument("--width", type=int, required=True)
args = parser.parse_args()
H = args.height
W = args.width
MODEL = args.model
model = onnx.load(f'{MODEL}')
onnx_session = onnxruntime.InferenceSession(f'{MODEL}')
inputs = {}
for i in onnx_session.get_inputs():
inputs[i.name] = [i.shape[0], i.shape[1], H, W]
print('@@@@@@@@@@@@@@@@@@@@@ inputs')
pprint(inputs)
model_simp, check = simplify(model, input_shapes=inputs)
basename_without_ext = \
os.path.splitext(os.path.basename(MODEL))[0].split('_')[0] \
+ os.path.splitext(os.path.basename(MODEL))[0].split('_')[1] \
+ os.path.splitext(os.path.basename(MODEL))[0].split('_')[2]
import onnx
print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
f = opt.weights.replace('.pt', '.onnx') # filename
model.fuse() # only for ONNX
torch.onnx.export(
model,
img,
f,
verbose=False,
opset_version=12,
input_names=['images'],
output_names=['classes', 'boxes'] if y is None else ['output'])
# Checks
onnx_model = onnx.load(f) # load onnx model
onnx.checker.check_model(onnx_model) # check onnx model
print(onnx.helper.printable_graph(
onnx_model.graph)) # print a human readable model
print('ONNX export success, saved as %s' % f)
except Exception as e:
print('ONNX export failure: %s' % e)
# CoreML export
try:
import coremltools as ct
print('\nStarting CoreML export with coremltools %s...' %
ct.__version__)
# convert model from torchscript and apply pixel scaling as per detect.py
model = ct.convert(ts,
def detect(data_dict,
device='',
checkpoint_path='checkpoint.pt',
img_size=320,
out='output',
cfg_path='cfg/yolov3-custom.cfg',
half=True,
view_img=False,
save_txt=False,
save_img=True,
conf_thres=0.3,
nms_thres=0.5):
img_size = (
320, 192
) if ONNX_EXPORT else img_size # (320, 192) or (416, 256) or (608, 352) for (height, width)
this_path = os.path.join(os.getcwd(), 'zoo/yolov3')
config_path = os.path.join(this_path, cfg_path)
source = data_dict['predict_on']
# Initialize
device = torch_utils.select_device(device='cpu' if ONNX_EXPORT else device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Initialize model
model = Darknet(config_path, img_size)
# Load checkpoint_path
checkpoint = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(checkpoint['model'])
# Eval mode
model.to(device).eval()
# Export mode
if ONNX_EXPORT:
img = torch.zeros((1, 3) + img_size) # (1, 3, 320, 192)
torch.onnx.export(model,
img,
'checkpoint_path/export.onnx',
verbose=False,
opset_version=11)
# Validate exported model
import onnx
model = onnx.load('checkpoint_path/export.onnx') # Load the ONNX model
onnx.checker.check_model(model) # Check that the IR is well formed
print(onnx.helper.printable_graph(
model.graph)) # Print a human readable representation of the graph
return
# Half precision
half = half and device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Set Dataloader
vid_path, vid_writer = None, None
dataset = LoadImages(source, img_size=img_size, half=half)
# Get classes and colors
classes = load_classes(data_dict['names'])
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(classes))]
# Run inference
t0 = time.time()
for path, img, im0s, vid_cap in dataset:
t = time.time()
# Get detections
img = torch.from_numpy(img).to(device)
if img.ndimension() == 3:
img = img.unsqueeze(0)
pred = model(img)[0]
if half:
pred = pred.float()
# Apply NMS
pred = non_max_suppression(pred, conf_thres, nms_thres)
# Process detections
for i, det in enumerate(pred): # detections per image
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
s += '%gx%g ' % img.shape[2:] # print string
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, classes[int(c)]) # add to string
# Write results
for *xyxy, conf, _, cls in det:
if save_txt: # Write to file
with open(save_path + '.txt', 'a') as file:
file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf))
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (classes[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)])
print('%sDone. (%.3fs)' % (s, time.time() - t))
# Stream results
if view_img:
cv2.imshow(p, im0)
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
if save_txt or save_img:
print('Results saved to %s' % out)
if platform == 'darwin': # MacOS
os.system('open ' + out + ' ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
