@@ -20,11 +20,59 @@ namespace tensorrt {
20
20
21
21
class Conv2dOpConverter : public OpConverter {
22
22
public:
23
- Conv2dOpConverter () {}
24
23
void operator ()(const framework::proto::OpDesc& op,
25
24
const framework::Scope& scope, bool test_mode) override {
26
25
LOG (INFO)
27
26
<< " convert a fluid conv2d op to tensorrt conv layer without bias" ;
27
+
28
+ framework::OpDesc op_desc (op, nullptr );
29
+ PADDLE_ENFORCE_EQ (op_desc.Input (" Input" ).size (), 1 );
30
+ PADDLE_ENFORCE_EQ (op_desc.Input (" Filter" ).size (), 1 ); // Y is a weight
31
+ PADDLE_ENFORCE_EQ (op_desc.Output (" Output" ).size (), 1 );
32
+
33
+ auto * X = engine_->GetITensor (op_desc.Input (" Input" ).front ());
34
+ // Declare weights
35
+ auto * Y_v = scope.FindVar (op_desc.Input (" Filter" ).front ());
36
+ PADDLE_ENFORCE_NOT_NULL (Y_v);
37
+ auto * Y_t = Y_v->GetMutable <framework::LoDTensor>();
38
+ auto * weight_data = Y_t->mutable_data <float >(platform::CPUPlace ());
39
+
40
+ const int n_output = Y_t->dims ()[0 ];
41
+ const int filter_h = Y_t->dims ()[2 ];
42
+ const int filter_w = Y_t->dims ()[3 ];
43
+
44
+ const int groups = boost::get<int >(op_desc.GetAttr (" groups" ));
45
+ const std::vector<int > dilations =
46
+ boost::get<std::vector<int >>(op_desc.GetAttr (" dilations" ));
47
+ const std::vector<int > strides =
48
+ boost::get<std::vector<int >>(op_desc.GetAttr (" strides" ));
49
+ const std::vector<int > paddings =
50
+ boost::get<std::vector<int >>(op_desc.GetAttr (" paddings" ));
51
+
52
+ nvinfer1::DimsHW nv_ksize (filter_h, filter_w);
53
+ nvinfer1::DimsHW nv_dilations (dilations[0 ], dilations[1 ]);
54
+ nvinfer1::DimsHW nv_strides (strides[0 ], strides[1 ]);
55
+ nvinfer1::DimsHW nv_paddings (paddings[0 ], paddings[1 ]);
56
+
57
+ TensorRTEngine::Weight weight{nvinfer1::DataType::kFLOAT ,
58
+ static_cast <void *>(weight_data),
59
+ Y_t->memory_size () / sizeof (float )};
60
+
61
+ TensorRTEngine::Weight bias{nvinfer1::DataType::kFLOAT , nullptr , 0 };
62
+ auto * layer = TRT_ENGINE_ADD_LAYER (
63
+ engine_, Convolution, *const_cast <nvinfer1::ITensor*>(X), n_output,
64
+ nv_ksize, weight.get (), bias.get ());
65
+ PADDLE_ENFORCE (layer != nullptr );
66
+ layer->setStride (nv_strides);
67
+ layer->setPadding (nv_paddings);
68
+ layer->setDilation (nv_dilations);
69
+ layer->setNbGroups (groups);
70
+
71
+ auto output_name = op_desc.Output (" Output" ).front ();
72
+ engine_->SetITensor (output_name, layer->getOutput (0 ));
73
+ if (test_mode) {
74
+ engine_->DeclareOutput (output_name);
75
+ }
28
76
}
29
77
};
30
78
0 commit comments