未验证 提交 de992400 编写于 作者: X xiaoxiaohehe001 提交者: GitHub

[Paddle Inference] Support dyanmic_shape of shuffle_channel. (#43488)

* shuffle_channel

* shuffle

* shuffle

* shuffle
上级 35694c2a
......@@ -39,25 +39,60 @@ class ShuffleChannelOpConverter : public OpConverter {
// Declare inputs
auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
auto input_dims = input->getDimensions();
int c = input_dims.d[0];
int h = input_dims.d[1];
int w = input_dims.d[2];
auto output_name = op_desc.Output("Out")[0];
int group = BOOST_GET_CONST(int, op_desc.GetAttr("group"));
auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
nvinfer1::Dims4 reshape_dim(group, c / group, h, w);
layer->setReshapeDimensions(reshape_dim);
layer->setSecondTranspose({1, 0, 2, 3});
auto* output = layer->getOutput(0);
#if IS_TRT_VERSION_GE(8000)
if (engine_->with_dynamic_shape()) {
auto* input_shape_tensor = Shape(input);
auto* channel_shape_tensor = GetEleTensorOfShape(input_shape_tensor, 1);
auto* group_tensor =
Add1DConstantLayer(group, output_name + "_group_tensor_");
auto* new_channel_shape_tensor = Div(channel_shape_tensor, group_tensor);
std::vector<int32_t> shape_dim3{0, 2, 3};
auto* shape_dim3_tensor = Gather(input_shape_tensor, shape_dim3);
auto* reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *output);
nvinfer1::Dims3 reshape_dim2(c, h, w);
reshape_layer->setReshapeDimensions(reshape_dim2);
std::vector<nvinfer1::ITensor*> itensors;
itensors.push_back(shape_dim3_tensor);
itensors.push_back(group_tensor);
itensors.push_back(new_channel_shape_tensor);
auto* reshape_tensor = Concat(itensors);
auto output_name = op_desc.Output("Out")[0];
RreplenishLayerAndOutput(reshape_layer, "shuffle_channel", {output_name},
test_mode);
auto* reshape_layer =
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *reshape_tensor);
nvinfer1::Permutation transpose_new_input{0, 3, 4, 1, 2};
reshape_layer->setSecondTranspose(transpose_new_input);
auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
layer->setInput(1, *(reshape_layer->getOutput(0)));
nvinfer1::Permutation transpose_embed{0, 2, 1, 3, 4};
layer->setSecondTranspose(transpose_embed);
auto* output = layer->getOutput(0);
auto* output_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *output);
output_layer->setInput(1, *input_shape_tensor);
RreplenishLayerAndOutput(output_layer, "shuffle_channel", {output_name},
test_mode);
}
#endif
if (!engine_->with_dynamic_shape()) {
int c = input_dims.d[0];
int h = input_dims.d[1];
int w = input_dims.d[2];
auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
nvinfer1::Dims4 reshape_dim(group, c / group, h, w);
layer->setReshapeDimensions(reshape_dim);
layer->setSecondTranspose({1, 0, 2, 3});
auto* output = layer->getOutput(0);
auto* reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *output);
nvinfer1::Dims3 reshape_dim2(c, h, w);
reshape_layer->setReshapeDimensions(reshape_dim2);
RreplenishLayerAndOutput(reshape_layer, "shuffle_channel", {output_name},
test_mode);
}
}
};
......
......@@ -1604,11 +1604,14 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8,
}
if (op_type == "shuffle_channel") {
#if !IS_TRT_VERSION_GE(8000)
if (with_dynamic_shape) {
VLOG(3) << "You are running the TRT Dynamic Shape mode, "
"the shuffle_channel op does not support dynamic shape yet";
"the shuffle_channel op does not support dynamic shape "
"trt versions below 8.0 yet";
return false;
}
#endif
}
if (op_type == "skip_layernorm") {
......
......@@ -77,7 +77,9 @@ class TrtConvertShuffleChannelTest(TrtLayerAutoScanTest):
self.dynamic_shape.opt_input_shape = {}
def generate_trt_nodes_num(attrs, dynamic_shape):
if dynamic_shape == True:
ver = paddle_infer.get_trt_compile_version()
if ver[0] * 1000 + ver[1] * 100 + ver[
2] * 10 < 8000 and dynamic_shape == True:
return 0, 3
else:
return 1, 2
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册