@@ -77,6 +77,39 @@ 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
+ // [1, 2]
86
+ float raw_bias[2 ] = {1.3 , 2.4 };
87
+
88
+ TensorRTEngine::Weight weight (nvinfer1::DataType::kFLOAT , raw_weight, 4 );
89
+ TensorRTEngine::Weight bias (nvinfer1::DataType::kFLOAT , raw_bias, 2 );
90
+ auto * x = engine_->DeclareInput (" x" , nvinfer1::DataType::kFLOAT ,
91
+ nvinfer1::DimsCHW{1 , 2 , 1 });
92
+ auto * fc_layer = TRT_ENGINE_ADD_LAYER (engine_, FullyConnected, *x, 2 ,
93
+ weight.get (), bias.get ());
94
+ PADDLE_ENFORCE (fc_layer != nullptr );
95
+
96
+ engine_->DeclareOutput (fc_layer, 0 , " y" );
97
+ engine_->FreezeNetwork ();
98
+ ASSERT_EQ (engine_->engine ()->getNbBindings (), 2 );
99
+
100
+ // fill in real data [1.0, 2.0]
101
+ float x_v[2 ] = {1.0 , 2.0 };
102
+ engine_->SetInputFromCPU (" x" , reinterpret_cast <void *>(&x_v),
103
+ 2 * sizeof (float ));
104
+ engine_->Execute (1 );
105
+
106
+ LOG (INFO) << " to get output" ;
107
+ float y_cpu[2 ] = {-1 ., -1 .};
108
+ engine_->GetOutputInCPU (" y" , &y_cpu[0 ], sizeof (float ) * 2 );
109
+ ASSERT_EQ (y_cpu[0 ], 4.5 );
110
+ ASSERT_EQ (y_cpu[1 ], 14.5 );
111
+ }
112
+
80
113
} // namespace tensorrt
81
114
} // namespace inference
82
115
} // namespace paddle
0 commit comments