diff --git a/paddle/fluid/inference/tensorrt/convert/softmax_op.cc b/paddle/fluid/inference/tensorrt/convert/softmax_op.cc index f401d407f78672735d474c7d56726cee30d5d599..624a1552b1813ef894f4a7a39624b495a7bea5aa 100644 --- a/paddle/fluid/inference/tensorrt/convert/softmax_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/softmax_op.cc @@ -33,26 +33,39 @@ class SoftMaxOpConverter : public OpConverter { auto* input1 = engine_->GetITensor(op_desc.Input("X")[0]); nvinfer1::Dims input_shape = input1->getDimensions(); int input_dims = input_shape.nbDims; - int axis = BOOST_GET_CONST(int, op_desc.GetAttr("axis")); + int axis = op_desc.HasAttr("axis") + ? BOOST_GET_CONST(int, op_desc.GetAttr("axis")) + : -1; auto* layer = TRT_ENGINE_ADD_LAYER(engine_, SoftMax, *const_cast(input1)); uint32_t axes = std::max(0, input_dims - 3); + // TODO(cryoco): Poor workaround. Fix padded dims problem when TRT layers + // support Nd. + int padded_dims = 0; + int explicit_batch = 0; + if (engine_->with_dynamic_shape()) explicit_batch = 1; + for (int i = input_dims - 1; i > explicit_batch; i--) { + if (input_shape.d[i] == 1) { + padded_dims += 1; + } else { + break; + } + } if (!engine_->with_dynamic_shape()) { if (axis == -1) { - axes = input_dims - 1; + axes = input_dims - 1 - padded_dims; } else { axes = axis; } - layer->setAxes(1 << axes); } else { if (axis == -1) { - axes = input_dims - 1; + axes = input_dims - 1 - padded_dims; } else { axes = axis + 1; } - layer->setAxes(1 << axes); } + layer->setAxes(1 << axes); auto output_name = op_desc.Output("Out")[0]; RreplenishLayerAndOutput(layer, "softmax", {output_name}, test_mode); diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc index 017afb5d7310a201c05dfcbe3c49a94cd89a7e92..3a8250815322b78c22c6ab7575ed098501fdfacb 100644 --- a/paddle/fluid/inference/tensorrt/op_teller.cc +++ b/paddle/fluid/inference/tensorrt/op_teller.cc @@ -107,8 +107,11 @@ bool OpTeller::Tell(const std::string& op_type, const framework::OpDesc& desc, op_type == "depthwise_conv2d" || op_type == "conv2d_transpose") { std::vector paddings = BOOST_GET_CONST(std::vector, desc.GetAttr("paddings")); - std::string padding_algorithm = - BOOST_GET_CONST(std::string, desc.GetAttr("padding_algorithm")); + + std::string padding_algorithm = "EXPLICIT"; + if (desc.HasAttr("padding_algorithm")) + padding_algorithm = + BOOST_GET_CONST(std::string, desc.GetAttr("padding_algorithm")); if (paddings.size() > 2 || (padding_algorithm == "SAME" && op_type != "pool2d")) return false;