| |
| |
| |
| |
| |
| from caffe2.proto import caffe2_pb2 |
| from caffe2.python import core, workspace |
| import onnx |
| import onnx.defs |
| from onnx.helper import make_node, make_graph, make_tensor_value_info, make_model |
| from onnx.backend.base import namedtupledict |
| from caffe2.python.models.download import ModelDownloader |
| import caffe2.python.onnx.backend as c2 |
| from caffe2.python.onnx.workspace import Workspace |
| from caffe2.python.trt.transform import convert_onnx_model_to_trt_op, transform_caffe2_net |
| from caffe2.python.onnx.tests.test_utils import TestCase |
| import numpy as np |
| import os.path |
| import time |
| import unittest |
| import tarfile |
| import tempfile |
| import shutil |
| from six.moves.urllib.request import urlretrieve |
| |
| def _print_net(net): |
| for i in net.external_input: |
| print("Input: {}".format(i)) |
| for i in net.external_output: |
| print("Output: {}".format(i)) |
| for op in net.op: |
| print("Op {}".format(op.type)) |
| for x in op.input: |
| print(" input: {}".format(x)) |
| for y in op.output: |
| print(" output: {}".format(y)) |
| |
| |
| def _base_url(opset_version): |
| return 'https://s3.amazonaws.com/download.onnx/models/opset_{}'.format(opset_version) |
| |
| # TODO: This is copied from https://github.com/onnx/onnx/blob/master/onnx/backend/test/runner/__init__.py. Maybe we should |
| # expose a model retrival API from ONNX |
| def _download_onnx_model(model_name, opset_version): |
| onnx_home = os.path.expanduser(os.getenv('ONNX_HOME', os.path.join('~', '.onnx'))) |
| models_dir = os.getenv('ONNX_MODELS', |
| os.path.join(onnx_home, 'models')) |
| model_dir = os.path.join(models_dir, model_name) |
| if not os.path.exists(os.path.join(model_dir, 'model.onnx')): |
| if os.path.exists(model_dir): |
| bi = 0 |
| while True: |
| dest = '{}.old.{}'.format(model_dir, bi) |
| if os.path.exists(dest): |
| bi += 1 |
| continue |
| shutil.move(model_dir, dest) |
| break |
| os.makedirs(model_dir) |
| |
| # On Windows, NamedTemporaryFile can not be opened for a |
| # second time |
| url = '{}/{}.tar.gz'.format(_base_url(opset_version), model_name) |
| download_file = tempfile.NamedTemporaryFile(delete=False) |
| try: |
| download_file.close() |
| print('Start downloading model {} from {}'.format( |
| model_name, url)) |
| urlretrieve(url, download_file.name) |
| print('Done') |
| with tarfile.open(download_file.name) as t: |
| t.extractall(models_dir) |
| except Exception as e: |
| print('Failed to prepare data for model {}: {}'.format( |
| model_name, e)) |
| raise |
| finally: |
| os.remove(download_file.name) |
| return model_dir |
| |
| class TensorRTOpTest(TestCase): |
| def setUp(self): |
| self.opset_version = onnx.defs.onnx_opset_version() |
| |
| def _test_relu_graph(self, X, batch_size, trt_max_batch_size): |
| node_def = make_node("Relu", ["X"], ["Y"]) |
| Y_c2 = c2.run_node(node_def, {"X": X}) |
| graph_def = make_graph( |
| [node_def], |
| name="test", |
| inputs=[make_tensor_value_info("X", onnx.TensorProto.FLOAT, [batch_size, 1, 3, 2])], |
| outputs=[make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [batch_size, 1, 3, 2])]) |
| model_def = make_model(graph_def, producer_name='relu-test') |
| op_outputs = [x.name for x in model_def.graph.output] |
| op = convert_onnx_model_to_trt_op(model_def, max_batch_size=trt_max_batch_size) |
| 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("X", X) |
| ws.RunOperatorsOnce([op]) |
| output_values = [ws.FetchBlob(name) for name in op_outputs] |
| Y_trt = namedtupledict('Outputs', op_outputs)(*output_values) |
| np.testing.assert_almost_equal(Y_c2, Y_trt) |
| |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_relu_graph_simple(self): |
| X = np.random.randn(1, 1, 3, 2).astype(np.float32) |
| self._test_relu_graph(X, 1, 50) |
| |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_relu_graph_big_batch(self): |
| X = np.random.randn(52, 1, 3, 2).astype(np.float32) |
| self._test_relu_graph(X, 52, 50) |
| |
| def _test_onnx_importer(self, model_name, data_input_index, opset_version=onnx.defs.onnx_opset_version()): |
| model_dir = _download_onnx_model(model_name, opset_version) |
| 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) |
| if opset_version >= 5: |
| # Some newer models from ONNX Zoo come with pre-set "data_0" input |
| ws.FeedBlob("data_0", 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) |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_resnet50(self): |
| self._test_onnx_importer('resnet50', 0, 9) |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_bvlc_alexnet(self): |
| self._test_onnx_importer('bvlc_alexnet', 0, 9) |
| |
| @unittest.skip("Until fixing Unsqueeze op") |
| def test_densenet121(self): |
| self._test_onnx_importer('densenet121', -1, 3) |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_inception_v1(self): |
| self._test_onnx_importer('inception_v1', -3, 9) |
| |
| @unittest.skip("Until fixing Unsqueeze op") |
| def test_inception_v2(self): |
| self._test_onnx_importer('inception_v2', 0, 9) |
| |
| @unittest.skip('Need to revisit our ChannelShuffle exporter to avoid generating 5D tensor') |
| def test_shufflenet(self): |
| self._test_onnx_importer('shufflenet', 0) |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_squeezenet(self): |
| self._test_onnx_importer('squeezenet', -1, 9) |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_vgg16(self): |
| self._test_onnx_importer('vgg16', 0, 9) |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_vgg19(self): |
| self._test_onnx_importer('vgg19', -2, 9) |
| |
| |
| class TensorRTTransformTest(TestCase): |
| def setUp(self): |
| self.model_downloader = ModelDownloader() |
| |
| def _add_head_tail(self, pred_net, new_head, new_tail): |
| orig_head = pred_net.external_input[0] |
| orig_tail = pred_net.external_output[0] |
| |
| # Add head |
| head = caffe2_pb2.OperatorDef() |
| head.type = "Copy" |
| head.input.append(new_head) |
| head.output.append(orig_head) |
| dummy = caffe2_pb2.NetDef() |
| dummy.op.extend(pred_net.op) |
| del pred_net.op[:] |
| pred_net.op.extend([head]) |
| pred_net.op.extend(dummy.op) |
| pred_net.external_input[0] = new_head |
| |
| # Add tail |
| tail = caffe2_pb2.OperatorDef() |
| tail.type = "Copy" |
| tail.input.append(orig_tail) |
| tail.output.append(new_tail) |
| pred_net.op.extend([tail]) |
| pred_net.external_output[0] = new_tail |
| |
| @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support") |
| def test_resnet50_core(self): |
| N = 2 |
| warmup = 20 |
| repeat = 100 |
| print("Batch size: {}, repeat inference {} times, warmup {} times".format(N, repeat, warmup)) |
| init_net, pred_net, _ = self.model_downloader.get_c2_model('resnet50') |
| self._add_head_tail(pred_net, 'real_data', 'real_softmax') |
| input_blob_dims = (N, 3, 224, 224) |
| input_name = "real_data" |
| |
| device_option = core.DeviceOption(caffe2_pb2.CUDA, 0) |
| init_net.device_option.CopyFrom(device_option) |
| pred_net.device_option.CopyFrom(device_option) |
| for op in pred_net.op: |
| op.device_option.CopyFrom(device_option) |
| op.engine = 'CUDNN' |
| net_outputs = pred_net.external_output |
| Y_c2 = None |
| data = np.random.randn(*input_blob_dims).astype(np.float32) |
| c2_time = 1 |
| workspace.SwitchWorkspace("gpu_test", True) |
| with core.DeviceScope(device_option): |
| workspace.FeedBlob(input_name, data) |
| workspace.RunNetOnce(init_net) |
| workspace.CreateNet(pred_net) |
| for _ in range(warmup): |
| workspace.RunNet(pred_net.name) |
| start = time.time() |
| for _ in range(repeat): |
| workspace.RunNet(pred_net.name) |
| end = time.time() |
| c2_time = end - start |
| output_values = [workspace.FetchBlob(name) for name in net_outputs] |
| Y_c2 = namedtupledict('Outputs', net_outputs)(*output_values) |
| workspace.ResetWorkspace() |
| |
| # Fill the workspace with the weights |
| with core.DeviceScope(device_option): |
| workspace.RunNetOnce(init_net) |
| |
| # Cut the graph |
| start = time.time() |
| pred_net_cut = transform_caffe2_net(pred_net, |
| {input_name: input_blob_dims}, |
| build_serializable_op=False) |
| del init_net, pred_net |
| pred_net_cut.device_option.CopyFrom(device_option) |
| for op in pred_net_cut.op: |
| op.device_option.CopyFrom(device_option) |
| #_print_net(pred_net_cut) |
| |
| Y_trt = None |
| input_name = pred_net_cut.external_input[0] |
| print("C2 runtime: {}s".format(c2_time)) |
| with core.DeviceScope(device_option): |
| workspace.FeedBlob(input_name, data) |
| workspace.CreateNet(pred_net_cut) |
| end = time.time() |
| print("Conversion time: {:.2f}s".format(end -start)) |
| |
| for _ in range(warmup): |
| workspace.RunNet(pred_net_cut.name) |
| start = time.time() |
| for _ in range(repeat): |
| workspace.RunNet(pred_net_cut.name) |
| end = time.time() |
| trt_time = end - start |
| print("TRT runtime: {}s, improvement: {}%".format(trt_time, (c2_time-trt_time)/c2_time*100)) |
| output_values = [workspace.FetchBlob(name) for name in net_outputs] |
| Y_trt = namedtupledict('Outputs', net_outputs)(*output_values) |
| np.testing.assert_allclose(Y_c2, Y_trt, rtol=1e-3) |
