fix batchnorm when inpu dims < 3 (#31933)

* fix batchnorm when inpu dims < 3

* add unittest for batchnorm dims = 2

paddle/fluid/inference/tensorrt/convert/batch_norm_op.cc

@@ -158,9 +158,29 @@ class BatchNormOpConverter : public OpConverter {
     TensorRTEngine::Weight power_weights{nvinfer1::DataType::kFLOAT, nullptr,
                                          0};
 
-    nvinfer1::IScaleLayer* layer =
-        TRT_ENGINE_ADD_LAYER(engine_, Scale, *const_cast<nvinfer1::ITensor*>(X),
-                             nvinfer1::ScaleMode::kCHANNEL, shift_weights.get(),
+    int dynamic_shape_offset = engine_->with_dynamic_shape() ? 1 : 0;
+    nvinfer1::ILayer* layer = nullptr;
+    nvinfer1::IShuffleLayer* expand_layer = nullptr;
+    nvinfer1::IShuffleLayer* squeeze_layer = nullptr;
+
+    auto x_dim = X->getDimensions();
+    if (x_dim.nbDims < 3 + dynamic_shape_offset) {
+      nvinfer1::Dims expand_shape;
+      expand_shape.nbDims = 3 + dynamic_shape_offset;
+      for (int i = 0; i < 3 + dynamic_shape_offset; i++) {
+        if (i < x_dim.nbDims) {
+          expand_shape.d[i] = x_dim.d[i] < 0 ? 0 : x_dim.d[i];
+        } else {
+          expand_shape.d[i] = 1;
+        }
+      }
+      expand_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *X);
+      expand_layer->setReshapeDimensions(expand_shape);
+      X = expand_layer->getOutput(0);
+    }
+
+    layer = TRT_ENGINE_ADD_LAYER(
+        engine_, Scale, *X, nvinfer1::ScaleMode::kCHANNEL, shift_weights.get(),
         scale_weights.get(), power_weights.get());
 
     auto output_name = op_desc.Output("Y").front();
@@ -168,7 +188,19 @@ class BatchNormOpConverter : public OpConverter {
                         std::move(combile_bias_tensor));
     engine_->SetWeights(op_desc.Input("Scale").front(),
                         std::move(combile_scale_tensor));
-    RreplenishLayerAndOutput(layer, "pool2d", {output_name}, test_mode);
+    if (x_dim.nbDims < 3 + dynamic_shape_offset) {
+      nvinfer1::Dims squeeze_shape;
+      squeeze_shape.nbDims = x_dim.nbDims;
+      for (int i = 0; i < squeeze_shape.nbDims; i++) {
+        squeeze_shape.d[i] = x_dim.d[i] < 0 ? 0 : x_dim.d[i];
+      }
+      squeeze_layer =
+          TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *(layer->getOutput(0)));
+      squeeze_layer->setReshapeDimensions(squeeze_shape);
+      layer = static_cast<nvinfer1::ILayer*>(squeeze_layer);
+    }
+    RreplenishLayerAndOutput(layer, "batchnorm_add_scale", {output_name},
+                             test_mode);
   }
 };
 

python/paddle/fluid/tests/unittests/ir/inference/test_trt_scale_op.py

@@ -48,5 +48,33 @@ class TRTScaleTest(InferencePassTest):
                 PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))
 
 
+class TRTScaleShape2Test(InferencePassTest):
+    def setUp(self):
+        with fluid.program_guard(self.main_program, self.startup_program):
+            data = fluid.data(
+                name="data", shape=[-1, 512, 512], dtype="float32")
+            scale_out = self.append_scale(data)
+            out = fluid.layers.batch_norm(scale_out, is_test=True)
+
+        self.feeds = {
+            "data": np.random.random([1, 512, 512]).astype("float32"),
+        }
+        self.enable_trt = True
+        self.trt_parameters = TRTScaleShape2Test.TensorRTParam(
+            1 << 30, 32, 1, AnalysisConfig.Precision.Float32, False, False)
+        self.fetch_list = [out]
+
+    def append_scale(self, data):
+        return fluid.layers.scale(
+            x=data, scale=2.0, bias=-1.0, bias_after_scale=False)
+
+    def test_check_output(self):
+        if core.is_compiled_with_cuda():
+            use_gpu = True
+            self.check_output_with_option(use_gpu, flatten=True)
+            self.assertTrue(
+                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))
+
+
 if __name__ == "__main__":
     unittest.main()