MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / parsePooling

Function parsePooling

parsers/caffe/caffeParser/opParsers/parsePooling.cpp:24–78  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

22namespace nvcaffeparser1
23{
24ILayer* parsePooling(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
25{
26 if (!checkBlobs(msg, 1, 1))
27 {
28 return nullptr;
29 }
30
31 const trtcaffe::PoolingParameter& p = msg.pooling_param();
32 if (p.pool() != trtcaffe::PoolingParameter::MAX && p.pool() != trtcaffe::PoolingParameter::AVE)
33 {
34 std::cout << "Caffe Parser: only AVE and MAX pool operations are supported" << std::endl;
35 return nullptr;
36 }
37
38 int kernelH, kernelW;
39 if (p.has_global_pooling() && p.global_pooling())
40 {
41 Dims3 dims = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions());
42 kernelH = dims.d[1];
43 kernelW = dims.d[2];
44 }
45 else
46 {
47 // mandatory
48 kernelH = p.has_kernel_h() ? p.kernel_h() : p.kernel_size();
49 kernelW = p.has_kernel_w() ? p.kernel_w() : p.kernel_size();
50 }
51
52 PoolingType type = p.has_pool() && p.pool() == trtcaffe::PoolingParameter::AVE ? PoolingType::kAVERAGE : PoolingType::kMAX;
53 auto layer = network.addPooling(*tensors[msg.bottom(0)], type, DimsHW{kernelH, kernelW});
54
55 if (layer)
56 {
57 int stride = p.has_stride() ? p.stride() : 1;
58 layer->setStride(DimsHW{p.has_stride_h() ? int(p.stride_h()) : stride, p.has_stride_w() ? int(p.stride_w()) : stride});
59
60 int pad = p.has_pad() ? p.pad() : 0;
61 layer->setPadding(DimsHW{p.has_pad_h() ? int(p.pad_h()) : pad, p.has_pad_w() ? int(p.pad_w()) : pad});
62
63 layer->setName(msg.name().c_str());
64 layer->setPaddingMode(PaddingMode::kCAFFE_ROUND_UP); // caffe pool use ceil mode by default
65 // FB pooling parameters
66 // Use floor((height + 2 * padding - kernel) / stride) + 1
67 // instead of ceil((height + 2 * padding - kernel) / stride) + 1
68 if (p.has_torch_pooling() ? p.torch_pooling() : false)
69 {
70 layer->setPaddingMode(PaddingMode::kCAFFE_ROUND_DOWN); // facebook torch pool use floor mode
71 }
72
73 tensors[msg.top(0)] = layer->getOutput(0);
74
75 layer->setAverageCountExcludesPadding(false); // unlike other frameworks, caffe use inclusive counting for padded averaging
76 }
77 return layer;
78}
79} //namespace nvcaffeparser1

Callers

nothing calls this directly

Calls 12

checkBlobsFunction · 0.85
getCHWFunction · 0.85
addPoolingMethod · 0.80
c_strMethod · 0.80
nameMethod · 0.80
getDimensionsMethod · 0.45
setStrideMethod · 0.45
setPaddingMethod · 0.45
setNameMethod · 0.45
setPaddingModeMethod · 0.45
getOutputMethod · 0.45

Tested by

no test coverage detected