if [[ $BUILD_ENVIRONMENT == *mkl* ]]; then
CMAKE_ARGS+=("-DBLAS=MKL")
fi
+
+if [[ $BUILD_ENVIRONMENT == py2-cuda9.0-cudnn7-ubuntu16.04 ]]; then
+
+ # removing http:// duplicate in favor of nvidia-ml.list
+ # which is https:// version of the same repo
+ sudo rm -f /etc/apt/sources.list.d/nvidia-machine-learning.list
+ curl -o ./nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb
+ sudo dpkg -i ./nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb
+ sudo apt-key add /var/nvinfer-runtime-trt-repo-5.0.2-ga-cuda9.0/7fa2af80.pub
+ sudo apt-get -qq update
+ sudo apt-get install libnvinfer5 libnvinfer-dev
+ rm ./nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb
+
+ CMAKE_ARGS+=("-DUSE_TENSORRT=ON")
+fi
+
if [[ $BUILD_ENVIRONMENT == *cuda* ]]; then
CMAKE_ARGS+=("-DUSE_CUDA=ON")
CMAKE_ARGS+=("-DCUDA_ARCH_NAME=Maxwell")
--- /dev/null
+# Caffe2 & TensorRT integration
+
+[](https://ci.pytorch.org/jenkins/job/caffe2-master)
+
+This directory contains the code implementing `TensorRTOp` Caffe2 operator as well as Caffe2 model converter (using `ONNX` model as an intermediate format).
+To enable this functionality in your PyTorch build please set
+
+`USE_TENSORRT=1 ... python setup.py ...`
+
+ or if you use CMake directly
+
+ `-DUSE_TENSORRT=ON`
+
+For further information please explore `caffe2/python/trt/test_trt.py` test showing all possible use cases.
+
+## Questions and Feedback
+
+Please use Github issues (https://github.com/pytorch/pytorch/issues) to ask questions, report bugs, and request new features.
// Otherwise, return the product of CHW dimensions
int64_t CheckDims(
const nvinfer1::Dims& nv_dims,
- const std::vector<int64_t>& c2_dims) {
+ at::ArrayRef<int64_t> c2_dims) {
if (nv_dims.nbDims + 1 != c2_dims.size()) {
CAFFE_THROW(
"Mismatched dimensions between TRT input (",
logger_(
(nvinfer1::ILogger::Severity)(OperatorBase::GetSingleArgument<int>(
"log_verbosity",
- FLAGS_minloglevel))),
+ FLAGS_caffe2_log_level))),
max_batch_size_(
OperatorBase::GetSingleArgument<int>("max_batch_size", 1)) {
{
CPUTensorToTensorProto(cpu_tensor, t);
} else if (BlobIsTensorType(*blob, CUDA)) {
const auto& cuda_tensor = blob->template Get<TensorCUDA>();
- const auto cpu_tensor = TensorCPU(cuda_tensor, context);
+ const auto cpu_tensor = TensorCPU(cuda_tensor, CPU);
context->FinishDeviceComputation();
CPUTensorToTensorProto(cpu_tensor, t);
} else {
model->set_producer_name("caffe2");
auto* opset_id = model->add_opset_import();
opset_id->set_domain("");
- opset_id->set_version(3);
+ opset_id->set_version(7);
}
} // namespace
::ONNX_NAMESPACE::TensorProto tf;
tf.set_name(t.name());
tf.mutable_dims()->CopyFrom(t.dims());
- tf.set_data_type(::ONNX_NAMESPACE::TensorProto::FLOAT);
- std::vector<int64_t> v;
- v.resize(t.raw_data().size() / sizeof(int64_t));
- memcpy(v.data(), t.raw_data().data(), t.raw_data().size());
- std::vector<float> vf;
- for (auto i : v) {
- vf.push_back(static_cast<float>(i));
- }
- tf.mutable_raw_data()->assign(
- reinterpret_cast<const char*>(vf.data()), sizeof(float) * vf.size());
+ if (t.data_type() == ::ONNX_NAMESPACE::TensorProto::FLOAT) {
+ tf.set_data_type(::ONNX_NAMESPACE::TensorProto::FLOAT);
+ std::vector<int64_t> v;
+ v.resize(t.raw_data().size() / sizeof(int64_t));
+ memcpy(v.data(), t.raw_data().data(), t.raw_data().size());
+ std::vector<float> vf;
+ for (auto i : v) {
+ vf.push_back(static_cast<float>(i));
+ }
+ tf.mutable_raw_data()->assign(
+ reinterpret_cast<const char *>(vf.data()), sizeof(float) * vf.size());
+ } else if (t.data_type() == ::ONNX_NAMESPACE::TensorProto::INT64) {
+ tf.set_data_type(::ONNX_NAMESPACE::TensorProto::INT64);
+ tf.mutable_raw_data()->assign(t.raw_data().data(), t.raw_data().size());
+ } else {
+ CAFFE_THROW("Unsupported tensor data type for conversion: ",
+ t.data_type());
+ }
onnx_model.mutable_graph()->add_initializer()->CopyFrom(tf);
}
}
auto shape_hints = InferShapes(&mapped_ws, pred_net, &shape_hints_ordered);
CAFFE_ENFORCE(pred_net, "Predict net cannot be nullptr");
- onnx::OnnxExporter exporter(nullptr, true);
+ onnx::OnnxExporter exporter(nullptr);
tensorrt::TrtLogger logger;
auto trt_builder = tensorrt::TrtObject(nvinfer1::createInferBuilder(logger));
auto trt_network = tensorrt::TrtObject(trt_builder->createNetwork());
// but it should be OK as the cost is really small. We also need to keep the
// same exporter throughout the process to avoid duplicated dummy name
// generation
- onnx::OnnxExporter exporter2(nullptr, true);
+ onnx::OnnxExporter exporter2(nullptr);
auto trt_converter = [this, &mapped_ws, &shape_hints, &exporter2](
const caffe2::NetDef& net) mutable {
return SubnetToTrtOp(net, &mapped_ws, &exporter2, &shape_hints);
print(" output: {}".format(y))
-_BASE_URL = 'https://s3.amazonaws.com/download.onnx/models/opset_{}'.format(onnx.defs.onnx_opset_version())
+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):
+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'))
# On Windows, NamedTemporaryFile can not be opened for a
# second time
- url = '{}/{}.tar.gz'.format(_BASE_URL, model_name)
+ url = '{}/{}.tar.gz'.format(_base_url(opset_version), model_name)
download_file = tempfile.NamedTemporaryFile(delete=False)
try:
download_file.close()
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 = 0):
- model_dir = _download_onnx_model(model_name)
+ 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]
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")
+ @unittest.skip("Until fixing Reshape op")
def test_resnet50(self):
- self._test_onnx_importer('resnet50')
+ self._test_onnx_importer('resnet50', 0)
@unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
def test_bvlc_alexnet(self):
- self._test_onnx_importer('bvlc_alexnet')
+ self._test_onnx_importer('bvlc_alexnet', 0)
- @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
+ @unittest.skip("Until fixing Unsqueeze op")
def test_densenet121(self):
- self._test_onnx_importer('densenet121', -1)
+ self._test_onnx_importer('densenet121', -1, 3)
- @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
+ @unittest.skip("Until fixing Reshape op")
def test_inception_v1(self):
- self._test_onnx_importer('inception_v1', -1)
+ self._test_onnx_importer('inception_v1', -1, 3)
- @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
+ @unittest.skip("Until fixing Reshape op")
def test_inception_v2(self):
- self._test_onnx_importer('inception_v2')
+ self._test_onnx_importer('inception_v2', 0, 3)
@unittest.skip('Need to revisit our ChannelShuffle exporter to avoid generating 5D tensor')
def test_shufflenet(self):
- self._test_onnx_importer('shufflenet')
+ self._test_onnx_importer('shufflenet', 0)
@unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
def test_squeezenet(self):
@unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
def test_vgg16(self):
- self._test_onnx_importer('vgg16')
+ self._test_onnx_importer('vgg16', 0)
- @unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
+ @unittest.skip("Until fixing Reshape op")
def test_vgg19(self):
- self._test_onnx_importer('vgg19', -1)
+ self._test_onnx_importer('vgg19', -1, 3)
class TensorRTTransformTest(DownloadingTestCase):
if (CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
if (USE_TENSORRT)
set(CMAKE_CUDA_COMPILER ${CUDA_NVCC_EXECUTABLE})
- add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/../third_party/onnx-tensorrt)
+ add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/../third_party/onnx-tensorrt EXCLUDE_FROM_ALL)
include_directories("${CMAKE_CURRENT_LIST_DIR}/../third_party/onnx-tensorrt")
caffe2_interface_library(nvonnxparser_static onnx_trt_library)
list(APPEND Caffe2_DEPENDENCY_WHOLE_LINK_LIBS onnx_trt_library)
# Before we run the setup_helpers, let's look for NO_* and WITH_*
# variables and hotpatch environment with the USE_* equivalent
use_env_vars = ['CUDA', 'CUDNN', 'FBGEMM', 'MIOPEN', 'MKLDNN', 'NNPACK', 'DISTRIBUTED',
- 'OPENCV', 'QNNPACK', 'FFMPEG', 'SYSTEM_NCCL', 'GLOO_IBVERBS']
+ 'OPENCV', 'TENSORRT', 'QNNPACK', 'FFMPEG', 'SYSTEM_NCCL',
+ 'GLOO_IBVERBS']
list(map(hotpatch_var, use_env_vars))
# Also hotpatch a few with BUILD_* equivalent
from tools.setup_helpers.cuda import USE_CUDA, CUDA_HOME, CUDA_VERSION
from tools.setup_helpers.build import (BUILD_BINARY, BUILD_TEST,
BUILD_CAFFE2_OPS, USE_LEVELDB,
- USE_LMDB, USE_OPENCV, USE_FFMPEG)
+ USE_LMDB, USE_OPENCV, USE_TENSORRT, USE_FFMPEG)
from tools.setup_helpers.rocm import USE_ROCM, ROCM_HOME, ROCM_VERSION
from tools.setup_helpers.cudnn import (USE_CUDNN, CUDNN_LIBRARY,
CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR)
my_env["USE_LEVELDB"] = "ON" if USE_LEVELDB else "OFF"
my_env["USE_LMDB"] = "ON" if USE_LMDB else "OFF"
my_env["USE_OPENCV"] = "ON" if USE_OPENCV else "OFF"
+ my_env["USE_TENSORRT"] = "ON" if USE_TENSORRT else "OFF"
my_env["USE_FFMPEG"] = "ON" if USE_FFMPEG else "OFF"
my_env["USE_DISTRIBUTED"] = "ON" if USE_DISTRIBUTED else "OFF"
my_env["USE_SYSTEM_NCCL"] = "ON" if USE_SYSTEM_NCCL else "OFF"
-Subproject commit fa0964e8477fc004ee2f49ee77ffce0bf7f711a9
+Subproject commit f1c7aa63d88d8d8ef70490f2ebb6b33f7450218b
-DUSE_LMDB=$USE_LMDB \
-DUSE_OPENCV=$USE_OPENCV \
-DUSE_QNNPACK=$USE_QNNPACK \
+ -DUSE_TENSORRT=$USE_TENSORRT \
-DUSE_FFMPEG=$USE_FFMPEG \
-DUSE_GLOG=OFF \
-DUSE_GFLAGS=OFF \
USE_LEVELDB = check_env_flag('USE_LEVELDB')
USE_LMDB = check_env_flag('USE_LMDB')
USE_OPENCV = check_env_flag('USE_OPENCV')
+USE_TENSORRT = check_env_flag('USE_TENSORRT')
USE_FFMPEG = check_env_flag('USE_FFMPEG')
projects / platform / upstream / pytorch.git / commitdiff