Skip to content

Commit 303277f

Browse files
authored
Merge pull request #10437 from panyx0718/infer2
Add a multi-dim add layer test.
2 parents 5ebf29b + 6728d96 commit 303277f

File tree

1 file changed

+31
-0
lines changed

1 file changed

+31
-0
lines changed

paddle/fluid/inference/tensorrt/test_engine.cc

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,37 @@ TEST_F(TensorRTEngineTest, add_layer) {
7777
ASSERT_EQ(y_cpu, x_v * 2 + 3);
7878
}
7979

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+
80111
} // namespace tensorrt
81112
} // namespace inference
82113
} // namespace paddle

0 commit comments

Comments
 (0)