Trained neural network using DIGITS is not working in DRIVE PX2


I am using NVIDIA DIGITS (v6) for training a custom multi class object detector. I have used the prototxt file from NVIDIA github ( and changed the parameters for rmulti class object detection.

The training was successful and I can see two objects being recognized in DIGITS. I got the tensorRT bin file and I deployed the network in DRIVE PX2. I have modified the sample code in PX2 for multiple class object detection.

Apparently, while doing inference on the video stream, I notices only one class is being detected.

While debugging, I noticed dwObjectDetector_initialize() needs to have following

a) The number of channels in the coverage blob is equal to the number of classes that the network detects.

b) The number of channels in the bounding box blob is equal to the 4 times number of classes that the network detects.

In my case, my coverage blob have 2 channels(for two classes), but my bounding box have just 4 channels.
After digging in prototxt file again, I notices the the bbox/regressor layer in caffe model needs to have 4 times the number of classes as an output(Eg, if I have 2 classes, I should have 8 outputs), but I have only 4 as the output.

If I try to change the parameter in prototxt file before training, it will gives an shape error in the next following layers.

Anyone knows the solution or how to implement the trained neural network in DRIVE PX2 for custom multi class object detector?

Thanks in advance