Skip to content

Commit 824b555

Browse files
bowang007bowang007
committed
Merge branch 'master' of https://github.com/NVIDIA/TRTorch into bowa_fallback
2 parents 3a72dc3 + 3589730 commit 824b555

File tree

11 files changed

+473
-5
lines changed

11 files changed

+473
-5
lines changed

‎core/conversion/converters/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ cc_library(
4444
"impl/matrix_multiply.cpp",
4545
"impl/pooling.cpp",
4646
"impl/reduce.cpp",
47+
"impl/replication_pad.cpp",
4748
"impl/shuffle.cpp",
4849
"impl/softmax.cpp",
4950
"impl/unary.cpp",

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

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
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 trtorch {
9+
namespace core {
10+
namespace conversion {
11+
namespace converters {
12+
namespace impl {
13+
namespace {
14+
15+
bool replication_padXd(ConversionCtx* ctx, const torch::jit::Node* n, args& args, int x_dim) {
16+
auto in = args[0].ITensor();
17+
auto inDims = in->getDimensions();
18+
int64_t inRank = inDims.nbDims;
19+
auto padding = args[1].unwrapToIntList().vec();
20+
if (padding.size() == 1) {
21+
for (int64_t i = 0; i < x_dim * 2 - 1; i++)
22+
padding.push_back(padding[0]);
23+
}
24+
if (inRank == 3) {
25+
TRTORCH_CHECK(padding.size() == 2, "3D tensors expect 2 values for padding");
26+
} else if (inRank == 4) {
27+
TRTORCH_CHECK(padding.size() == 4, "4D tensors expect 4 values for padding");
28+
} else if (inRank == 5) {
29+
TRTORCH_CHECK(padding.size() == 6, "5D tensors expect 6 values for padding");
30+
} else {
31+
TRTORCH_THROW_ERROR("Only 3D, 4D, 5D padding with non-constant padding are supported for now");
32+
}
33+
34+
std::vector<nvinfer1::ITensor*> tensors_vec;
35+
// input: (N, C, D_in, H_in, W_in).
36+
// padding: (padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back)
37+
// When axis is inRank - 1, making W_out = W_in + padding_left + padding_right.
38+
// When axis is inRank - 2, making H_out = H_in + padding_top + padding_bottom.
39+
// When axis is inRank - 1, making D_out = D_in + padding_front + padding_back.
40+
for (int64_t i = 0; i < int(padding.size() / 2); i++) {
41+
int64_t axis = inRank - (i + 1); // axis = {inRank - 1, inRank - 2, inRank - 3}
42+
int64_t padding_index = i * 2;
43+
44+
if (padding[padding_index] > 0) { // left/top/front padding value
45+
tensors_vec.clear();
46+
at::Tensor left_indices = torch::tensor({0}, torch::kInt32);
47+
auto indicesTensor = tensor_to_const(ctx, left_indices);
48+
auto left_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
49+
auto left_gather_out = left_gather_layer->getOutput(0);
50+
for (int i = 0; i < padding[padding_index]; i++) {
51+
tensors_vec.push_back(left_gather_out);
52+
}
53+
tensors_vec.push_back(in);
54+
auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
55+
concat_layer->setAxis(axis);
56+
in = concat_layer->getOutput(0);
57+
inDims = in->getDimensions();
58+
}
59+
60+
if (padding[padding_index + 1] > 0) { // right/bottom/back padding value
61+
tensors_vec.clear();
62+
tensors_vec.push_back(in);
63+
64+
nvinfer1::ITensor* indicesTensor = NULL;
65+
if (inDims.d[axis] == -1) {
66+
auto shapeTensor = ctx->net->addShape(*in)->getOutput(0);
67+
at::Tensor dimValue = torch::tensor({axis}, torch::kInt32);
68+
auto dimTensor = tensor_to_const(ctx, dimValue);
69+
indicesTensor = ctx->net->addGather(*shapeTensor, *dimTensor, 0)->getOutput(0);
70+
} else {
71+
auto indices = torch::tensor({inDims.d[axis] - 1}, torch::kInt32);
72+
indicesTensor = tensor_to_const(ctx, indices);
73+
}
74+
auto right_gather_layer = ctx->net->addGather(*in, *indicesTensor, axis);
75+
auto right_gather_out = right_gather_layer->getOutput(0);
76+
77+
for (int i = 0; i < padding[padding_index + 1]; i++) {
78+
tensors_vec.push_back(right_gather_out);
79+
}
80+
81+
auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
82+
concat_layer->setAxis(axis);
83+
in = concat_layer->getOutput(0);
84+
inDims = in->getDimensions();
85+
}
86+
}
87+
88+
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
89+
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
90+
91+
return true;
92+
}
93+
94+
auto replication_pad_registrations TRTORCH_UNUSED =
95+
RegisterNodeConversionPatterns()
96+
.pattern({"aten::replication_pad1d(Tensor self, int[2] padding) -> (Tensor)",
97+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
98+
replication_padXd(ctx, n, args, 1);
99+
return true;
100+
}})
101+
.pattern({"aten::replication_pad2d(Tensor self, int[4] padding) -> (Tensor)",
102+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
103+
replication_padXd(ctx, n, args, 2);
104+
return true;
105+
}})
106+
.pattern({"aten::replication_pad3d(Tensor self, int[6] padding) -> (Tensor)",
107+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
108+
replication_padXd(ctx, n, args, 3);
109+
return true;
110+
}});
111+
112+
} // namespace
113+
} // namespace impl
114+
} // namespace converters
115+
} // namespace conversion
116+
} // namespace core
117+
} // namespace trtorch

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

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,81 @@ static auto shuffle_registrations TRTORCH_UNUSED =
125125
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
126126
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
127127

128+
return true;
129+
}})
130+
.pattern({"aten::pixel_shuffle(Tensor self, int upscale_factor) -> (Tensor)",
131+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
132+
auto self = args[0].ITensorOrFreeze(ctx);
133+
auto in_shape = util::toVec(self->getDimensions());
134+
int64_t irank = in_shape.size();
135+
TRTORCH_CHECK(
136+
irank >= 3,
137+
"pixel_shuffle expects input to have at least 3 dimensions, but got input with "
138+
<< irank << " dimension(s)");
139+
int64_t upscale_factor = args[1].unwrapToInt();
140+
TRTORCH_CHECK(
141+
upscale_factor > 0,
142+
"pixel_shuffle expects a positive upscale_factor, but got " << upscale_factor);
143+
int64_t upscale_factor_squared = upscale_factor * upscale_factor;
144+
145+
const auto NUM_NON_BATCH_DIMS = 3;
146+
const auto self_sizes_batch_end = in_shape.end() - NUM_NON_BATCH_DIMS;
147+
148+
int64_t ic = in_shape[irank - 3];
149+
int64_t ih = in_shape[irank - 2];
150+
int64_t iw = in_shape[irank - 1];
151+
152+
TRTORCH_CHECK(
153+
ic % upscale_factor_squared == 0,
154+
"pixel_shuffle expects its input's 'channel' dimension to be divisible by the square of "
155+
<< "upscale_factor, but input.size(-3)=" << ic << " is not divisible by "
156+
<< upscale_factor_squared);
157+
158+
int64_t oc = ic / upscale_factor_squared;
159+
int64_t oh = ih * upscale_factor;
160+
int64_t ow = iw * upscale_factor;
161+
162+
// First, reshape to split the channels dim from c into 3 separate dims: (oc,
163+
// upscale_factor, upscale_factor). This allows shuffling to be done next by
164+
// permuting dims.
165+
std::vector<int64_t> added_dims_shape(in_shape.begin(), self_sizes_batch_end);
166+
added_dims_shape.insert(added_dims_shape.end(), {oc, upscale_factor, upscale_factor, ih, iw});
167+
auto view_layer = ctx->net->addShuffle(*self);
168+
TRTORCH_CHECK(view_layer, "Unable to create shuffle layer from node: " << *n);
169+
view_layer->setReshapeDimensions(util::toDims(added_dims_shape));
170+
int64_t view_rank = added_dims_shape.size();
171+
172+
// Next, shuffle by permuting the new upscale_factor dims alongside the height and width dims.
173+
auto permutation_layer = ctx->net->addShuffle(*view_layer->getOutput(0));
174+
TRTORCH_CHECK(permutation_layer, "Unable to create shuffle layer from node: " << *n);
175+
// std::iota is used to maintain the batch dims within the permutation.
176+
// Eg: if added_dims_shape is {n1, n2, c, r, r, h, w}, then the new_order is {view_rank-7,
177+
// view_rank-6, view_rank-5, view_rank-2, view_rank-4, view_rank-1, view_rank-3}
178+
std::vector<int64_t> new_order(in_shape.begin(), self_sizes_batch_end);
179+
std::iota(new_order.begin(), new_order.end(), 0);
180+
new_order.insert(
181+
new_order.end(),
182+
{view_rank - 5 /* oc */,
183+
view_rank - 2 /* ih */,
184+
view_rank - 4 /* 1st upscale_factor */,
185+
view_rank - 1 /* iw */,
186+
view_rank - 3 /* 2nd upscale_factor */});
187+
nvinfer1::Permutation permute;
188+
std::copy(new_order.begin(), new_order.end(), permute.order);
189+
permutation_layer->setSecondTranspose(permute);
190+
191+
// Finally, upscale by collapsing (ih, upscale_factor) -> a single dim (oh)
192+
// and (iw, upscale_factor) -> a single dim (ow).
193+
std::vector<int64_t> final_shape(in_shape.begin(), self_sizes_batch_end);
194+
final_shape.insert(final_shape.end(), {oc, oh, ow});
195+
auto last_view_layer = ctx->net->addShuffle(*permutation_layer->getOutput(0));
196+
TRTORCH_CHECK(last_view_layer, "Unable to create shuffle layer from node: " << *n);
197+
last_view_layer->setReshapeDimensions(util::toDims(final_shape));
198+
last_view_layer->setName(util::node_info(n).c_str());
199+
200+
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], last_view_layer->getOutput(0));
201+
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
202+
128203
return true;
129204
}});
130205

‎docs/_notebooks/Resnet50-example.html

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -690,7 +690,7 @@
690690
</div>
691691
</div>
692692
<p>
693-
<img alt="f03b9ffb995e463a9ece48dbc03b19a6" src="http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png"/>
693+
<img alt="2f9e31058b32421bb811779059a6af66" src="http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png"/>
694694
</p>
695695
<h1 id="notebooks-resnet50-example--page-root">
696696
TRTorch Getting Started - ResNet 50

‎docs/_notebooks/lenet-getting-started.html

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -784,7 +784,7 @@
784784
</div>
785785
</div>
786786
<p>
787-
<img alt="cd6d780b542849ab939e734e03646533" src="http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png"/>
787+
<img alt="ab920573d9c6495cbe539fd91862cc7e" src="http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png"/>
788788
</p>
789789
<h1 id="notebooks-lenet-getting-started--page-root">
790790
TRTorch Getting Started - LeNet

‎docs/_notebooks/ssd-object-detection-demo.html

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -804,7 +804,7 @@
804804
</div>
805805
</div>
806806
<p>
807-
<img alt="01bd2524df734a8aabc6c06d402cfd04" src="http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png"/>
807+
<img alt="d216ad9778614725a54e3da2ff4c800a" src="http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png"/>
808808
</p>
809809
<h1 id="notebooks-ssd-object-detection-demo--page-root">
810810
Object Detection with TRTorch (SSD)

‎docs/py_api/trtorch.html

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1011,7 +1011,7 @@ <h2 id="functions">
10111011
<span class="sig-paren">
10121012
)
10131013
</span>
1014-
→ &lt;torch._C.ScriptClass object at 0x7fc84234b370&gt;
1014+
→ &lt;torch._C.ScriptClass object at 0x7fb1218be8f0&gt;
10151015
<a class="headerlink" href="#trtorch.TensorRTCompileSpec" title="Permalink to this definition">
10161016
10171017
</a>

‎docs/searchindex.js

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tests/core/conversion/converters/BUILD

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,10 @@ converter_test(
4747
name = "test_reduce"
4848
)
4949

50+
converter_test(
51+
name = "test_replication_pad"
52+
)
53+
5054
converter_test(
5155
name = "test_shuffle"
5256
)
@@ -99,6 +103,7 @@ test_suite(
99103
":test_matrix_multiply",
100104
":test_pooling",
101105
":test_reduce",
106+
":test_replication_pad",
102107
":test_shuffle",
103108
":test_softmax",
104109
":test_unary",

0 commit comments

Comments
 (0)