jkjung-avt / tensorrt_demos

TensorRT MODNet, YOLOv4, YOLOv3, SSD, MTCNN, and GoogLeNet
https://jkjung-avt.github.io/
MIT License
1.74k stars 545 forks source link

Is it possible to use CBN? #498

Closed PedroCarrinhoRibeiro closed 2 years ago

PedroCarrinhoRibeiro commented 2 years ago

Hi,

So in Darknet, it's possible to turn on Cross-iteration Batch Normalization by changing batch_normalize=1 to batch_normalize=2 in the convolutional layers in the .cfg file.

However, when I tried to convert a model that used CBN from Darknet to trt, the model couldn't make any inference.

Is CBN implemented in your repo? If so, what changes do I need to make to use it?

Thanks in advance.

Edit: P.S. The model works completely fine in trt if I use the regular batch normalization instead.

jkjung-avt commented 2 years ago

Thanks for pointing out the problem. It was because the code only handled batch_normalize=1 case (2 was not handled correctly).

https://github.com/jkjung-avt/tensorrt_demos/blob/ca6ac81b4c7558122df3ff3018479412828a88a4/yolo/yolo_to_onnx.py#L652

The problem has been fixed in the latest code (https://github.com/jkjung-avt/tensorrt_demos/commit/f53b5ae9b004489463a407d8e9b230f39230d051) now.