From 3de43a87ef1a94c0fc7161a922bcbeb078705e48 Mon Sep 17 00:00:00 2001 From: Xin Pan Date: Sun, 6 May 2018 22:02:28 +0800 Subject: [PATCH] Add a multi-dim add layer test. We need to figure out if tensorrt use row-major or col-major for tensor layerout inorder to do conversion. --- .../fluid/inference/tensorrt/test_engine.cc | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/paddle/fluid/inference/tensorrt/test_engine.cc b/paddle/fluid/inference/tensorrt/test_engine.cc index a08b78f93..e27d73314 100644 --- a/paddle/fluid/inference/tensorrt/test_engine.cc +++ b/paddle/fluid/inference/tensorrt/test_engine.cc @@ -77,6 +77,39 @@ TEST_F(TensorRTEngineTest, add_layer) { ASSERT_EQ(y_cpu, x_v * 2 + 3); } +TEST_F(TensorRTEngineTest, add_layer_multi_dim) { + // Weight in CPU memory. + // It seems tensorrt FC use col-major: [[1.0, 3.3], [1.1, 4.4]] + // instead of row-major, which is [[1.0, 1.1], [3.3, 4.4]] + float raw_weight[4] = {1.0, 1.1, 3.3, 4.4}; + // [1, 2] + float raw_bias[2] = {1.3, 2.4}; + + TensorRTEngine::Weight weight(nvinfer1::DataType::kFLOAT, raw_weight, 4); + TensorRTEngine::Weight bias(nvinfer1::DataType::kFLOAT, raw_bias, 2); + auto* x = engine_->DeclareInput("x", nvinfer1::DataType::kFLOAT, + nvinfer1::DimsCHW{1, 2, 1}); + auto* fc_layer = TRT_ENGINE_ADD_LAYER(engine_, FullyConnected, *x, 2, + weight.get(), bias.get()); + PADDLE_ENFORCE(fc_layer != nullptr); + + engine_->DeclareOutput(fc_layer, 0, "y"); + engine_->FreezeNetwork(); + ASSERT_EQ(engine_->engine()->getNbBindings(), 2); + + // fill in real data [1.0, 2.0] + float x_v[2] = {1.0, 2.0}; + engine_->SetInputFromCPU("x", reinterpret_cast(&x_v), + 2 * sizeof(float)); + engine_->Execute(1); + + LOG(INFO) << "to get output"; + float y_cpu[2] = {-1., -1.}; + engine_->GetOutputInCPU("y", &y_cpu[0], sizeof(float) * 2); + ASSERT_EQ(y_cpu[0], 4.5); + ASSERT_EQ(y_cpu[1], 14.5); +} + } // namespace tensorrt } // namespace inference } // namespace paddle -- GitLab