Skip to content

Commit 474e7bf

Browse files
bowang007bowang007
committed
Merge branch 'master' of https://github.com/NVIDIA/Torch-TensorRT into maskrcnn
2 parents ffaa4bf + ba9f730 commit 474e7bf

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+1709
-664
lines changed

‎README.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ Resources:
1212
- [Comprehensive Discusion (GTC Event)](https://www.nvidia.com/en-us/on-demand/session/gtcfall21-a31107/)
1313
- [Pre-built Docker Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch). To use this container, make an NGC account and sign in to NVIDIA's registry with an API key. Refer to [this guide](https://docs.nvidia.com/ngc/ngc-catalog-user-guide/index.html#registering-activating-ngc-account) for the same.
1414

15+
## NVIDIA NGC Container
16+
Torch-TensorRT is distributed in the ready-to-run NVIDIA [NGC PyTorch Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) starting with 21.11. We recommend using this prebuilt container to experiment & develop with Torch-TensorRT; it has all dependencies with the proper versions as well as example notebooks included.
1517

1618
## Building a docker container for Torch-TensorRT
1719

@@ -48,7 +50,7 @@ If you would like to build outside a docker container, please follow the section
4850
// Set input datatypes. Allowerd options torch::{kFloat, kHalf, kChar, kInt32, kBool}
4951
// Size of input_dtypes should match number of inputs to the network.
5052
// If input_dtypes is not set, default precision follows traditional PyT / TRT rules
51-
auto input = torch_tensorrt::Input(dims, torch::kHalf)
53+
auto input = torch_tensorrt::Input(dims, torch::kHalf);
5254
auto compile_settings = torch_tensorrt::ts::CompileSpec({input});
5355
// FP16 execution
5456
compile_settings.enabled_precisions = {torch::kHalf};

‎WORKSPACE

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -128,18 +128,18 @@ http_archive(
128128
# build_file = "@//third_party/tensorrt/local:BUILD"
129129
#)
130130

131-
#########################################################################
132-
# Testing Dependencies (optional - comment out on aarch64)
133-
#########################################################################
134-
pip_install(
135-
name = "torch_tensorrt_py_deps",
136-
requirements = "//py:requirements.txt",
137-
)
138-
139-
pip_install(
140-
name = "py_test_deps",
141-
requirements = "//tests/py:requirements.txt",
142-
)
131+
# #########################################################################
132+
# # Testing Dependencies (optional - comment out on aarch64)
133+
# #########################################################################
134+
# pip_install(
135+
# name = "torch_tensorrt_py_deps",
136+
# requirements = "//py:requirements.txt",
137+
# )
138+
139+
# pip_install(
140+
# name = "py_test_deps",
141+
# requirements = "//tests/py:requirements.txt",
142+
# )
143143

144144
pip_install(
145145
name = "pylinter_deps",

‎core/conversion/conversion.cpp

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,8 @@ void AddLayer(ConversionCtx* ctx, const torch::jit::Node* n) {
105105
// Node input has not been converted yet or is a prim op
106106
TORCHTRT_THROW_ERROR(
107107
"Unable to retrieve all node inputs for node: "
108-
<< util::node_info(n) << " (ctx.AddLayer)\nSpecifically failed to retrieve value for input: " << *input_node);
108+
<< util::node_info(n) << " (ctx.AddLayer)\nSpecifically failed to retrieve value for input: %"
109+
<< input->debugName());
109110
}
110111
}
111112

@@ -534,18 +535,22 @@ bool VerifyConverterSupportForBlock(const torch::jit::Block* b, bool suppress_er
534535
if (unsupported_ops.size() != 0) {
535536
std::stringstream unsupported_msg;
536537
unsupported_msg
537-
<< "Method requested cannot be compiled by Torch-TensorRT.TorchScript.\nUnsupported operators listed below:"
538+
<< "Method requested cannot be compiled end to end by Torch-TensorRT.TorchScript.\nUnsupported operators listed below:"
538539
<< std::endl;
539540
for (auto s : unsupported_ops) {
540541
unsupported_msg << " - " << s.second << std::endl;
541542
}
542-
unsupported_msg << "You can either implement converters for these ops in your application or request implementation"
543-
<< std::endl;
544-
unsupported_msg << "https://www.github.com/nvidia/Torch-TensorRT/issues" << std::endl;
545-
unsupported_msg << std::endl << "In Module:" << std::endl;
546543

547544
if (!suppress_errors) {
545+
unsupported_msg
546+
<< "You can either implement converters for these ops in your application or request implementation"
547+
<< std::endl;
548+
unsupported_msg << "https://www.github.com/nvidia/Torch-TensorRT/issues" << std::endl;
549+
unsupported_msg << std::endl << "In Module:" << std::endl;
550+
548551
LOG_ERROR(unsupported_msg.str());
552+
} else {
553+
LOG_INFO(unsupported_msg.str());
549554
}
550555

551556
std::unordered_map<std::string, std::unordered_set<std::string>> unsupported_node_locations;
@@ -572,7 +577,11 @@ bool VerifyConverterSupportForBlock(const torch::jit::Block* b, bool suppress_er
572577
traceback << str;
573578
}
574579
auto tb_str = traceback.str();
575-
LOG_ERROR(tb_str);
580+
if (!suppress_errors) {
581+
LOG_ERROR(tb_str);
582+
} else {
583+
LOG_DEBUG(tb_str);
584+
}
576585
}
577586

578587
return false;

‎core/conversion/converters/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ cc_library(
7171
"impl/pooling.cpp",
7272
"impl/quantization.cpp",
7373
"impl/reduce.cpp",
74+
"impl/reflection_pad.cpp",
7475
"impl/replication_pad.cpp",
7576
"impl/select.cpp",
7677
"impl/shuffle.cpp",

‎core/conversion/converters/impl/reflection_pad.cpp

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
#include <ATen/ATen.h>
2+
#include <vector>
3+
#include "NvInfer.h"
4+
#include "core/conversion/converters/converters.h"
5+
#include "core/util/prelude.h"
6+
#include "torch/torch.h"
7+
8+
namespace torch_tensorrt {
9+
namespace core {
10+
namespace conversion {
11+
namespace converters {
12+
namespace impl {
13+
namespace {
14+
15+
auto reflection_padXd TORCHTRT_UNUSED =
16+
RegisterNodeConversionPatterns()
17+
.pattern({"aten::reflection_pad2d(Tensor self, int[4] padding) -> (Tensor)",
18+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
19+
auto in = args[0].ITensor();
20+
auto inDims = in->getDimensions();
21+
int64_t inRank = inDims.nbDims;
22+
auto padding = args[1].unwrapToIntList().vec();
23+
if (padding.size() == 1) {
24+
for (int64_t i = 0; i < 3; i++)
25+
padding.push_back(padding[0]);
26+
}
27+
if (inRank == 4) {
28+
TORCHTRT_CHECK(padding.size() == 4, "4D tensors expect 4 values for padding");
29+
} else {
30+
TORCHTRT_THROW_ERROR("Only 4D padding are supported for now");
31+
}
32+
33+
std::vector<nvinfer1::ITensor*> tensors_vec;
34+
// 2d padding: (padding_left, padding_right, padding_top, padding_bottom)
35+
36+
for (int64_t i = 0; i < int(padding.size() / 2); i++) {
37+
int64_t axis = inRank - (i + 1); // axis = {inRank - 1, inRank - 2}
38+
int64_t padding_index = i * 2;
39+
40+
if (padding[padding_index] > 0) { // left/top padding value
41+
tensors_vec.clear();
42+
43+
for (int i = 0; i < padding[padding_index]; i++) {
44+
at::Tensor left_indices = torch::tensor({padding[padding_index] - i}, torch::kInt32);
45+
auto indicesTensor = tensor_to_const(ctx, left_indices);
46+
auto left_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
47+
auto left_gather_out = left_gather_layer->getOutput(0);
48+
tensors_vec.push_back(left_gather_out);
49+
}
50+
tensors_vec.push_back(in);
51+
auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
52+
concat_layer->setAxis(axis);
53+
in = concat_layer->getOutput(0);
54+
inDims = in->getDimensions();
55+
}
56+
57+
if (padding[padding_index + 1] > 0) { // right/bottom padding value
58+
tensors_vec.clear();
59+
tensors_vec.push_back(in);
60+
61+
for (int i = 0; i < padding[padding_index + 1]; i++) {
62+
nvinfer1::ITensor* indicesTensor = NULL;
63+
auto indices = torch::tensor({inDims.d[axis] - 1 - (i + 1)}, torch::kInt32);
64+
indicesTensor = tensor_to_const(ctx, indices);
65+
auto right_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
66+
auto right_gather_out = right_gather_layer->getOutput(0);
67+
tensors_vec.push_back(right_gather_out);
68+
}
69+
70+
auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
71+
concat_layer->setAxis(axis);
72+
in = concat_layer->getOutput(0);
73+
inDims = in->getDimensions();
74+
}
75+
}
76+
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
77+
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
78+
79+
return true;
80+
}})
81+
.pattern({"aten::reflection_pad1d(Tensor self, int[2] padding) -> (Tensor)",
82+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
83+
auto in = args[0].ITensor();
84+
auto inDims = in->getDimensions();
85+
int64_t inRank = inDims.nbDims;
86+
auto padding = args[1].unwrapToIntList().vec();
87+
if (padding.size() == 1) {
88+
for (int64_t i = 0; i < 1; i++)
89+
padding.push_back(padding[0]);
90+
}
91+
92+
std::vector<nvinfer1::ITensor*> tensors_vec;
93+
// 1d padding: (padding_left, padding_right)
94+
95+
int64_t axis = inRank - 1;
96+
int64_t padding_index = 0;
97+
98+
if (padding[padding_index] > 0) { // left padding value
99+
tensors_vec.clear();
100+
101+
for (int i = 0; i < padding[padding_index]; i++) {
102+
at::Tensor left_indices = torch::tensor({padding[padding_index] - i}, torch::kInt32);
103+
auto indicesTensor = tensor_to_const(ctx, left_indices);
104+
auto left_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
105+
auto left_gather_out = left_gather_layer->getOutput(0);
106+
tensors_vec.push_back(left_gather_out);
107+
}
108+
tensors_vec.push_back(in);
109+
auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
110+
concat_layer->setAxis(axis);
111+
in = concat_layer->getOutput(0);
112+
inDims = in->getDimensions();
113+
}
114+
115+
if (padding[padding_index + 1] > 0) { // right padding value
116+
tensors_vec.clear();
117+
tensors_vec.push_back(in);
118+
119+
for (int i = 0; i < padding[padding_index + 1]; i++) {
120+
nvinfer1::ITensor* indicesTensor = NULL;
121+
auto indices = torch::tensor({inDims.d[axis] - 1 - (i + 1)}, torch::kInt32);
122+
indicesTensor = tensor_to_const(ctx, indices);
123+
auto right_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
124+
auto right_gather_out = right_gather_layer->getOutput(0);
125+
tensors_vec.push_back(right_gather_out);
126+
}
127+
128+
auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
129+
concat_layer->setAxis(axis);
130+
in = concat_layer->getOutput(0);
131+
inDims = in->getDimensions();
132+
}
133+
134+
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
135+
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
136+
137+
return true;
138+
}});
139+
140+
} // namespace
141+
} // namespace impl
142+
} // namespace converters
143+
} // namespace conversion
144+
} // namespace core
145+
} // namespace torch_tensorrt

‎core/conversion/evaluators/aten.cpp

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
#include <math.h>
2+
13
#include "ATen/core/List.h"
24
#include "ATen/core/functional.h"
35
#include "ATen/core/ivalue.h"
@@ -98,6 +100,17 @@ DEFINE_GENERIC_TWO_INPUT_EVALUATOR(
98100
"aten::ge.float_int(float a, int b) -> (bool)",
99101
}));
100102

103+
DEFINE_ARITHMATIC_TWO_INPUT_EVALUATOR(
104+
pow,
105+
"aten::pow",
106+
pow(a, b),
107+
std::set<std::string>({
108+
"aten::pow.int(int a, int b) -> (float)",
109+
"aten::pow.float(float a, float b) -> (float)",
110+
"aten::pow.int_float(int a, float b) -> (float)",
111+
"aten::pow.float_int(float a, int b) -> (float)",
112+
}));
113+
101114
DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(
102115
and,
103116
"aten::__and__",
@@ -285,6 +298,31 @@ auto aten_registrations TORCHTRT_UNUSED =
285298
EvalOptions().validSchemas({
286299
"aten::append.t(t[](a!) self, t(c -> *) el) -> (t[](a!))",
287300
})})
301+
.evaluator({c10::Symbol::fromQualString("aten::extend"),
302+
[](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
303+
if (args.at(n->input(0)).IValue()->isList() && args.at(n->input(1)).IValue()->isList()) {
304+
c10::IValue* self_ptr = args.at(n->input(0)).IValueMut();
305+
auto self = self_ptr->to<c10::List<c10::IValue>>();
306+
auto other = args.at(n->input(1)).IValue()->to<c10::List<c10::IValue>>();
307+
const int64_t other_size = other.size();
308+
309+
// Modify value in place
310+
for (int64_t i = 0; i < other_size; i++) {
311+
self.push_back(other.get(i));
312+
}
313+
314+
*self_ptr = c10::IValue(self);
315+
return {};
316+
} else {
317+
TORCHTRT_THROW_ERROR(
318+
"Unimplemented data type for aten::extend.t evaluator: "
319+
<< args.at(n->input(0)).IValue()->type()->str() << ", "
320+
<< args.at(n->input(1)).IValue()->type()->str());
321+
}
322+
},
323+
EvalOptions().validSchemas({
324+
"aten::extend.t(t[](a!) self, t[] other) -> ()",
325+
})})
288326
.evaluator({c10::Symbol::fromQualString("aten::neg"),
289327
[](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
290328
auto el = args.at(n->input(0)).unwrapToInt();

‎core/conversion/evaluators/eval_macros.h

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,53 @@
7777
}, \
7878
EvalOptions().validSchemas(schemas)});
7979

80+
#define DEFINE_ARITHMATIC_TWO_INPUT_EVALUATOR(name, node_kind, operation, schemas) \
81+
auto name##_registrations TORCHTRT_UNUSED = RegisterNodeEvaluators().evaluator( \
82+
{c10::Symbol::fromQualString(node_kind), \
83+
[](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> { \
84+
if (args.at(n->input(0)).IValue()->isInt()) { \
85+
auto a = args.at(n->input(0)).unwrapToInt(); \
86+
if (args.at(n->input(1)).IValue()->isInt()) { \
87+
auto b = args.at(n->input(1)).unwrapToInt(); \
88+
return operation; \
89+
} else if (args.at(n->input(1)).IValue()->isDouble()) { \
90+
auto b = args.at(n->input(1)).unwrapToDouble(); \
91+
return operation; \
92+
} else if (args.at(n->input(1)).IValue()->isBool()) { \
93+
auto b = args.at(n->input(1)).unwrapToBool(); \
94+
return operation; \
95+
} else { \
96+
TORCHTRT_THROW_ERROR( \
97+
"Unimplemented data type for " \
98+
<< node_kind << " evaluator b arg:" << args.at(n->input(1)).IValue()->type()->str()); \
99+
return {}; \
100+
} \
101+
} else if (args.at(n->input(0)).IValue()->isDouble()) { \
102+
auto a = args.at(n->input(0)).unwrapToDouble(); \
103+
if (args.at(n->input(1)).IValue()->isInt()) { \
104+
auto b = args.at(n->input(1)).unwrapToInt(); \
105+
return operation; \
106+
} else if (args.at(n->input(1)).IValue()->isDouble()) { \
107+
auto b = args.at(n->input(1)).unwrapToDouble(); \
108+
return operation; \
109+
} else if (args.at(n->input(1)).IValue()->isBool()) { \
110+
auto b = args.at(n->input(1)).unwrapToBool(); \
111+
return operation; \
112+
} else { \
113+
TORCHTRT_THROW_ERROR( \
114+
"Unimplemented data type for " \
115+
<< node_kind << " evaluator b arg:" << args.at(n->input(1)).IValue()->type()->str()); \
116+
return {}; \
117+
} \
118+
} else { \
119+
TORCHTRT_THROW_ERROR( \
120+
"Unimplemented data type for " \
121+
<< node_kind << " evaluator a arg: " << args.at(n->input(0)).IValue()->type()->str()); \
122+
return {}; \
123+
} \
124+
}, \
125+
EvalOptions().validSchemas(schemas)});
126+
80127
#define DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(node_kind, node_name, operation, type, schemas) \
81128
auto node_kind##_registrations TORCHTRT_UNUSED = RegisterNodeEvaluators().evaluator( \
82129
{c10::Symbol::fromQualString(node_name), \

0 commit comments

Comments
 (0)