search

LitLeo / TensorRT_Tutorial

981 stars 184 forks source link

如何处理网络模型中的BatchNorm层？ #3

Closed wuxler closed 6 years ago

wuxler commented 6 years ago

我想用TensorRT来加速YOLO v2模型，但是YOLO v2模型中包含有BatchNorm层，程序会在ICudaEngine engine = builder->buildCudaEngine(network)这一步中断报错，官方文档中指的Batch Normalization可以用Scale层代替，具体应该怎么做呢？

rencheng123456 commented 6 years ago

你好，我最近也在研究使用TensorRT来加速YOLO v2模型，想问下，但不知道如何着手。希望能给予我一些指导

jsmlay commented 6 years ago

TensorRT 3.0 Have YOLO plugins headers. Plz double check before usage. Find it in NvInferPlugin.h in /usr/src/TensorRT/include. It need Jetpack 3.2 support. When deploy it, you need to write plugin, plz refer to FasterRCNN sample. Sorry for using English because i'm on TX2.

LitLeo commented 6 years ago

抱歉最近没怎么关注git。BN的实现可以用scale层代替，直接将BN层的公式转换成scale层公式就可以了。我之前的实现代码如下，仅供参考。

                float *scale_values = new float[l_ptr->out_c];
                float *shift_values = new float[l_ptr->out_c];
                for (int i = 0; i < l_ptr->out_c; ++i)
                {
                    scale_values[i] = 1 / (sqrt(l_ptr->rolling_variance[i] + .00001f)) * l_ptr->scales[i];
                    shift_values[i] = -l_ptr->rolling_mean[i] / (sqrt(l_ptr->rolling_variance[i] + .00001f)) * l_ptr->scales[i] + l_ptr->biases[i];
                }
                Weights scale {DataType::kFLOAT, scale_values, l_ptr->out_c};
                Weights shift {DataType::kFLOAT, shift_values, l_ptr->out_c};
                Weights power {DataType::kFLOAT, nullptr, 0};

                auto scale_1 = network->addScale(*inputs_v[0], ScaleMode::kCHANNEL, shift, scale, power);
                assert(scale_1 != nullptr);
                inputs_v[0] = scale_1->getOutput(0);
wuxler commented 6 years ago

已经解决了，非常感谢！