@@ -24,19 +24,19 @@ class LayerNormOpConverter : public OpConverter {
24
24
void operator ()(const framework::proto::OpDesc& op,
25
25
const framework::Scope& scope,
26
26
bool test_mode) override {
27
- VLOG (4 ) << " convert a layer_norm op with dynamic shape to Normalization "
28
- " layer or Static shape tensorrt layer_norm plugin" ;
27
+ VLOG (4 ) << " convert a layer_norm op to INormalization layer or "
28
+ " layer_norm plugin" ;
29
29
framework::OpDesc op_desc (op, nullptr );
30
-
31
30
auto * X = engine_->GetITensor (op_desc.Input (" X" )[0 ]);
32
- auto rank = X->getDimensions ().nbDims ;
33
31
std::string output_name = op_desc.Output (" Y" )[0 ];
34
32
const float eps = op_desc.HasAttr (" epsilon" )
35
33
? PADDLE_GET_CONST (float , op_desc.GetAttr (" epsilon" ))
36
34
: 1e-5f ;
37
35
if (engine_->with_dynamic_shape ()) {
36
+ #if IS_TRT_VERSION_GE(8600)
38
37
auto * Scale = engine_->GetITensor (op_desc.Input (" Scale" )[0 ]);
39
38
auto * Bias = engine_->GetITensor (op_desc.Input (" Bias" )[0 ]);
39
+ auto rank = X->getDimensions ().nbDims ;
40
40
int32_t begin_axis =
41
41
op_desc.HasAttr (" begin_norm_axis" )
42
42
? PADDLE_GET_CONST (int , op_desc.GetAttr (" begin_norm_axis" ))
@@ -67,61 +67,54 @@ class LayerNormOpConverter : public OpConverter {
67
67
Scale,
68
68
concat_shape_tensor,
69
69
(" layer_norm Scale: reshape: (Output(" + output_name + " )" ).c_str ());
70
- #if IS_TRT_VERSION_GE(8600)
71
70
auto layer = TRT_ENGINE_ADD_LAYER (
72
71
engine_, Normalization, *X, *Scale_reshape, *Bias_reshape, axisMask);
73
72
layer->setEpsilon (eps);
74
73
RreplenishLayerAndOutput (layer, " layer_norm" , {output_name}, test_mode);
75
- #else
76
- // μ
77
- auto miu_layer = TRT_ENGINE_ADD_LAYER (
78
- engine_, Reduce, *X, nvinfer1::ReduceOperation::kAVG , axisMask, true );
79
- miu_layer->setName ((output_name + " _miu" ).c_str ());
80
- auto miu_output = miu_layer->getOutput (0 );
81
- // x−μ
82
- auto xsubmiu_output = Sub (X, miu_output);
83
- // σ
84
- // pow(x−μ,2)
85
- auto pow_tensor = Add1DConstantLayer (static_cast <float >(2 ));
86
- auto xsubmiu_pow_out = Pow (
87
- xsubmiu_output,
88
- BroadcastTensors (xsubmiu_output,
89
- pow_tensor,
90
- (" layer_norm_pow: reshape_for_broadcast: (Output(" +
91
- output_name + " )" )
92
- .c_str ()));
93
- // mean_var
94
- auto mean_var_layer =
95
- TRT_ENGINE_ADD_LAYER (engine_,
96
- Reduce,
97
- *xsubmiu_pow_out,
98
- nvinfer1::ReduceOperation::kAVG ,
99
- axisMask,
100
- true );
101
- mean_var_layer->setName ((output_name + " _sigma" ).c_str ());
102
- auto mean_var_out = mean_var_layer->getOutput (0 );
103
- // sigma
104
- auto eps_tensor = Add1DConstantLayer (eps);
105
- auto sum_out = Sum (
106
- mean_var_out,
107
- BroadcastTensors (mean_var_out,
108
- eps_tensor,
109
- (" layer_norm_eps: reshape_for_broadcast: (Output(" +
110
- output_name + " )" )
111
- .c_str ()));
112
- auto sigma_layer = TRT_ENGINE_ADD_LAYER (
113
- engine_, Unary, *sum_out, nvinfer1::UnaryOperation::kSQRT );
114
- auto sigma_output = sigma_layer->getOutput (0 );
115
- // σ/sigma
116
- auto div_out = Div (xsubmiu_output, sigma_output);
117
- // (σ/sigma)*g+b
118
- auto scale_out = Prod (div_out, Scale_reshape);
119
- auto layer = TRT_ENGINE_ADD_LAYER (engine_,
120
- ElementWise,
121
- *scale_out,
122
- *Bias_reshape,
123
- nvinfer1::ElementWiseOperation::kSUM );
124
- RreplenishLayerAndOutput (layer, " layer_norm" , {output_name}, test_mode);
74
+ #endif
75
+ #if IS_TRT_VERSION_LT(8600)
76
+ // For dynamic shape & trt<8.6,
77
+ // the shape of mean and variance will be determine in configuPlugin.
78
+ auto * X = engine_->GetITensor (op_desc.Input (" X" ).front ());
79
+ auto * Bias_v = scope.FindVar (op_desc.Input (" Bias" ).front ());
80
+ auto * Scale_v = scope.FindVar (op_desc.Input (" Scale" ).front ());
81
+ const int begin_norm_axis =
82
+ op_desc.HasAttr (" begin_norm_axis" )
83
+ ? PADDLE_GET_CONST (int , op_desc.GetAttr (" begin_norm_axis" ))
84
+ : 1 ;
85
+ PADDLE_ENFORCE_NOT_NULL (
86
+ Bias_v,
87
+ platform::errors::InvalidArgument (
88
+ " Input(Bias) of layer_norm should not be null." ));
89
+ PADDLE_ENFORCE_NOT_NULL (
90
+ Scale_v,
91
+ platform::errors::InvalidArgument (
92
+ " Input(Scale) of layer_norm should not be null." ));
93
+ auto * Bias_t = Bias_v->GetMutable <phi::DenseTensor>();
94
+ auto * Scale_t = Scale_v->GetMutable <phi::DenseTensor>();
95
+ auto bias_weight =
96
+ engine_->GetFp32TrtWeight (op_desc.Input (" Bias" ).front (), *Bias_t);
97
+ auto scale_weight =
98
+ engine_->GetFp32TrtWeight (op_desc.Input (" Scale" ).front (), *Scale_t);
99
+ nvinfer1::ILayer* layernorm_layer = nullptr ;
100
+ std::vector<int64_t > mean_shape{1 };
101
+ std::vector<int64_t > variance_shape{1 };
102
+ bool with_fp16 =
103
+ engine_->WithFp16 () && !engine_->disable_trt_plugin_fp16 ();
104
+ plugin::LayerNormPluginDynamic* plugin =
105
+ new plugin::LayerNormPluginDynamic (
106
+ static_cast <const float *>(bias_weight.get ().values ),
107
+ bias_weight.get ().count ,
108
+ static_cast <const float *>(scale_weight.get ().values ),
109
+ scale_weight.get ().count ,
110
+ begin_norm_axis,
111
+ eps,
112
+ mean_shape,
113
+ variance_shape,
114
+ with_fp16);
115
+ layernorm_layer = engine_->AddDynamicPlugin (&X, 1 , plugin);
116
+ RreplenishLayerAndOutput (
117
+ layernorm_layer, " layer_norm" , {output_name}, test_mode);
125
118
#endif
126
119
} else {
127
120
auto * Bias_v = scope.FindVar (op_desc.Input (" Bias" )[0 ]);
0 commit comments