Description
I have a simple neural network written in TensorFlow 1.15 that I need to convert to a TensorRT optimized plan using tensorRT_optimizer tool which is inside DriveWorks SDK 1.2 in DRIVE PX2. The tool requires that the input dimensions be CXHXW. But the input dimensions of the TensorFlow models always have the batch dimension as the first dimension. Therefore, always the input dimensions are in the shape of BxCxHxW. Is there a way to remove this batch dimension from the TensorFlow computational graph only for inference?
Relevant Screenshots
A part of the visualized computational graph of lenet5_custom_pool.uff which is provided in DriveWorks samples:
You can see that the input dimension is 1x28x28.
A part of the visualized computational graph of the same lenet5 network that I trained:
