@@ -25,12 +25,42 @@ TEST(OpConverter, ConvertBlock) {
25
25
framework::ProgramDesc prog;
26
26
auto * block = prog.MutableBlock (0 );
27
27
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
+
28
39
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" });
29
43
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
31
55
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*/ );
34
64
}
35
65
36
66
} // namespace tensorrt
0 commit comments