@@ -77,6 +77,37 @@ TEST_F(TensorRTEngineTest, add_layer) {
77
77
ASSERT_EQ (y_cpu, x_v * 2 + 3 );
78
78
}
79
79
80
+ TEST_F (TensorRTEngineTest, add_layer_multi_dim) {
81
+ // Weight in CPU memory.
82
+ // It seems tensorrt FC use col-major: [[1.0, 3.3], [1.1, 4.4]]
83
+ // instead of row-major, which is [[1.0, 1.1], [3.3, 4.4]]
84
+ float raw_weight[4 ] = {1.0 , 1.1 , 3.3 , 4.4 };
85
+ float raw_bias[2 ] = {1.3 , 2.4 };
86
+
87
+ TensorRTEngine::Weight weight (nvinfer1::DataType::kFLOAT , raw_weight, 4 );
88
+ TensorRTEngine::Weight bias (nvinfer1::DataType::kFLOAT , raw_bias, 2 );
89
+ auto * x = engine_->DeclareInput (" x" , nvinfer1::DataType::kFLOAT ,
90
+ nvinfer1::DimsCHW{1 , 2 , 1 });
91
+ auto * fc_layer = TRT_ENGINE_ADD_LAYER (engine_, FullyConnected, *x, 2 ,
92
+ weight.get (), bias.get ());
93
+ PADDLE_ENFORCE (fc_layer != nullptr );
94
+
95
+ engine_->DeclareOutput (fc_layer, 0 , " y" );
96
+ engine_->FreezeNetwork ();
97
+ ASSERT_EQ (engine_->engine ()->getNbBindings (), 2 );
98
+
99
+ float x_v[2 ] = {1.0 , 2.0 };
100
+ engine_->SetInputFromCPU (" x" , reinterpret_cast <void *>(&x_v),
101
+ 2 * sizeof (float ));
102
+ engine_->Execute (1 );
103
+
104
+ LOG (INFO) << " to get output" ;
105
+ float y_cpu[2 ] = {-1 ., -1 .};
106
+ engine_->GetOutputInCPU (" y" , &y_cpu[0 ], sizeof (float ) * 2 );
107
+ ASSERT_EQ (y_cpu[0 ], 4.5 );
108
+ ASSERT_EQ (y_cpu[1 ], 14.5 );
109
+ }
110
+
80
111
} // namespace tensorrt
81
112
} // namespace inference
82
113
} // namespace paddle
0 commit comments