Skip to content

Commit 0015df1

Browse files
committed
modify op converter for conv2d
1 parent a5c96af commit 0015df1

File tree

1 file changed

+33
-3
lines changed

1 file changed

+33
-3
lines changed

paddle/fluid/inference/tensorrt/convert/test_op_converter.cc

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,42 @@ TEST(OpConverter, ConvertBlock) {
2525
framework::ProgramDesc prog;
2626
auto* block = prog.MutableBlock(0);
2727
auto* conv2d_op = block->AppendOp();
28+
29+
// init trt engine
30+
cudaStream_t stream_;
31+
std::unique_ptr<TensorRTEngine> engine_;
32+
engine_.reset(new TensorRTEngine(5, 1 << 15, &stream_));
33+
engine_->InitNetwork();
34+
PADDLE_ENFORCE_EQ(cudaStreamCreate(&stream_), 0);
35+
36+
engine_->DeclareInput("conv2d-X", nvinfer1::DataType::kFLOAT,
37+
nvinfer1::Dims3(2, 5, 5));
38+
2839
conv2d_op->SetType("conv2d");
40+
conv2d_op->SetInput("Input", {"conv2d-X"});
41+
conv2d_op->SetInput("Filter", {"conv2d-Y"});
42+
conv2d_op->SetOutput("Output", {"conv2d-Out"});
2943

30-
OpConverter converter;
44+
const std::vector<int> strides({1, 1});
45+
const std::vector<int> paddings({1, 1});
46+
const std::vector<int> dilations({1, 1});
47+
const int groups = 1;
48+
49+
conv2d_op->SetAttr("strides", strides);
50+
conv2d_op->SetAttr("paddings", paddings);
51+
conv2d_op->SetAttr("dilations", dilations);
52+
conv2d_op->SetAttr("groups", groups);
53+
54+
// init scope
3155
framework::Scope scope;
32-
converter.ConvertBlock(*block->Proto(), {}, scope,
33-
nullptr /*TensorRTEngine*/);
56+
std::vector<int> dim_vec = {3, 2, 3, 3};
57+
auto* x = scope.Var("conv2d-Y");
58+
auto* x_tensor = x->GetMutable<framework::LoDTensor>();
59+
x_tensor->Resize(framework::make_ddim(dim_vec));
60+
61+
OpConverter converter;
62+
converter.ConvertBlock(*block->Proto(), {"conv2d-Y"}, scope,
63+
engine_.get() /*TensorRTEngine*/);
3464
}
3565

3666
} // namespace tensorrt

0 commit comments

Comments
 (0)