def test_tensorrt_inference():
"""Run LeNet-5 inference comparison between MXNet and TensorRT."""
check_tensorrt_installation()
mnist = mx.test_utils.get_mnist()
num_epochs = 10
batch_size = 128
model_name = 'lenet5'
model_dir = os.getenv("LENET_MODEL_DIR", "/tmp")
model_file = '%s/%s-symbol.json' % (model_dir, model_name)
params_file = '%s/%s-%04d.params' % (model_dir, model_name, num_epochs)
_, _, _, all_test_labels = get_iters(mnist, batch_size)
# Load serialized MXNet model (model-symbol.json + model-epoch.params)
sym, arg_params, aux_params = mx.model.load_checkpoint(model_name, num_epochs)
print("LeNet-5 test")
print("Running inference in MXNet")
mx_pct = run_inference(sym, arg_params, aux_params, mnist, all_test_labels,
batch_size=batch_size, use_tensorrt=False)
print("Running inference in MXNet-TensorRT")
trt_pct = run_inference(sym, arg_params, aux_params, mnist, all_test_labels,
batch_size=batch_size, use_tensorrt=True)
print("MXNet accuracy: %f" % mx_pct)
print("MXNet-TensorRT accuracy: %f" % trt_pct)
absolute_accuracy_diff = abs(mx_pct - trt_pct)
epsilon = 1.01e-2
assert absolute_accuracy_diff < epsilon, \
"""Absolute diff. between MXNet & TensorRT accuracy (%f) exceeds threshold (%f):
MXNet = %f, TensorRT = %f""" % (absolute_accuracy_diff, epsilon, mx_pct, trt_pct)
def test_tensorrt_inference():
"""Run LeNet-5 inference comparison between MXNet and TensorRT."""
original_try_value = mx.contrib.tensorrt.get_use_tensorrt()
try:
check_tensorrt_installation()
mnist = mx.test_utils.get_mnist()
num_epochs = 10
batch_size = 128
model_name = 'lenet5'
model_dir = os.getenv("LENET_MODEL_DIR", "/tmp")
model_file = '%s/%s-symbol.json' % (model_dir, model_name)
params_file = '%s/%s-%04d.params' % (model_dir, model_name, num_epochs)
_, _, _, all_test_labels = get_iters(mnist, batch_size)
# Load serialized MXNet model (model-symbol.json + model-epoch.params)
sym, arg_params, aux_params = mx.model.load_checkpoint(model_name, num_epochs)
print("LeNet-5 test")
print("Running inference in MXNet")
mx_pct = run_inference(sym, arg_params, aux_params, mnist, all_test_labels,
batch_size=batch_size, use_tensorrt=False)
print("Running inference in MXNet-TensorRT")
trt_pct = run_inference(sym, arg_params, aux_params, mnist, all_test_labels,
batch_size=batch_size, use_tensorrt=True)
print("MXNet accuracy: %f" % mx_pct)
print("MXNet-TensorRT accuracy: %f" % trt_pct)
assert abs(mx_pct - trt_pct) < 1e-2, \
"""Diff. between MXNet & TensorRT accuracy too high:
MXNet = %f, TensorRT = %f""" % (mx_pct, trt_pct)
finally:
mx.contrib.tensorrt.set_use_tensorrt(original_try_value)
def test_tensorrt_inference():
"""Run LeNet-5 inference comparison between MXNet and TensorRT."""
original_try_value = mx.contrib.tensorrt.get_use_tensorrt()
try:
check_tensorrt_installation()
mnist = mx.test_utils.get_mnist()
num_epochs = 10
batch_size = 128
model_name = 'lenet5'
model_dir = os.getenv("LENET_MODEL_DIR", "/tmp")
model_file = '%s/%s-symbol.json' % (model_dir, model_name)
params_file = '%s/%s-%04d.params' % (model_dir, model_name, num_epochs)
_, _, _, all_test_labels = get_iters(mnist, batch_size)
# Load serialized MXNet model (model-symbol.json + model-epoch.params)
sym, arg_params, aux_params = mx.model.load_checkpoint(
model_name, num_epochs)
print("LeNet-5 test")
print("Running inference in MXNet")
mx_pct = run_inference(sym,
arg_params,
aux_params,
mnist,
all_test_labels,
batch_size=batch_size,
use_tensorrt=False)
print("Running inference in MXNet-TensorRT")
trt_pct = run_inference(sym,
arg_params,
aux_params,
mnist,
all_test_labels,
batch_size=batch_size,
use_tensorrt=True)
print("MXNet accuracy: %f" % mx_pct)
print("MXNet-TensorRT accuracy: %f" % trt_pct)
assert abs(mx_pct - trt_pct) < 1e-2, \
"""Diff. between MXNet & TensorRT accuracy too high:
MXNet = %f, TensorRT = %f""" % (mx_pct, trt_pct)
finally:
mx.contrib.tensorrt.set_use_tensorrt(original_try_value)
optimizer='sgd',
optimizer_params={'learning_rate': 0.1, 'momentum': 0.9},
eval_metric='acc',
batch_end_callback=mx.callback.Speedometer(batch_size, 1),
num_epoch=num_epochs)
# predict accuracy for lenet
acc = mx.metric.Accuracy()
lenet_model.score(test_iter, acc)
accuracy = acc.get()[1]
assert accuracy > 0.95, "LeNet-5 training accuracy on MNIST was too low"
return lenet_model
if __name__ == '__main__':
num_epochs = 10
batch_size = 128
model_name = 'lenet5'
model_dir = os.getenv("LENET_MODEL_DIR", "/tmp")
model_file = '%s/%s-symbol.json' % (model_dir, model_name)
params_file = '%s/%s-%04d.params' % (model_dir, model_name, num_epochs)
if not (os.path.exists(model_file) and os.path.exists(params_file)):
mnist = mx.test_utils.get_mnist()
_, _, _, all_test_labels = get_iters(mnist, batch_size)
trained_lenet = train_lenet5(num_epochs, batch_size,
*get_iters(mnist, batch_size)[:-1])
trained_lenet.save_checkpoint(model_name, num_epochs)
'momentum': 0.9
},
eval_metric='acc',
batch_end_callback=mx.callback.Speedometer(batch_size, 1),
num_epoch=num_epochs)
# predict accuracy for lenet
acc = mx.metric.Accuracy()
lenet_model.score(test_iter, acc)
accuracy = acc.get()[1]
assert accuracy > 0.95, "LeNet-5 training accuracy on MNIST was too low"
return lenet_model
if __name__ == '__main__':
num_epochs = 10
batch_size = 128
model_name = 'lenet5'
model_dir = os.getenv("LENET_MODEL_DIR", "/tmp")
model_file = '%s/%s-symbol.json' % (model_dir, model_name)
params_file = '%s/%s-%04d.params' % (model_dir, model_name, num_epochs)
if not (os.path.exists(model_file) and os.path.exists(params_file)):
mnist = mx.test_utils.get_mnist()
_, _, _, all_test_labels = get_iters(mnist, batch_size)
trained_lenet = train_lenet5(num_epochs, batch_size,
*get_iters(mnist, batch_size)[:-1])
trained_lenet.save_checkpoint(model_name, num_epochs)