ceccocats / tkDNN

Deep neural network library and toolkit to do high performace inference on NVIDIA jetson platforms
GNU General Public License v2.0
717 stars 209 forks source link

Additional support for custom network architectures #176

Open nightduck opened 3 years ago

nightduck commented 3 years ago

This fixes 2 issues: one where NetworkRT hardcodes the names of input and output layers (issue #154), and another where it permits use of 2D layers, like fully-connected layers (issue #175)

nightduck commented 3 years ago

NetworkRT.cpp, as is, assumes the batch dimension to be 1. The yolo models in this repo follow that convention, and I have not editted them, lest this PR be too disruptive. But other TensorRT engines generated through other means may not keep to this convention. Or if they do, they may explicity declare the batch dimension (eg 1x3x224x224). Keras works this way. When NetworkRT loads in their dimensions, you get:

n: 1 c: 1 h: 3 w: 224 l: 224

A PR that says "adds suppport for custom network architectures" and still forces them to have an omitted batch dimension is misleading. This likely shouldn't be accepted until that is adressed.

nightduck commented 3 years ago

Added another commit that will, by default, assume the TensorRT engine does not specify the batch dimension. (This is the current default behaviour in master branch). However, it also allows the developer to specify that the batch dimension is present, or that the dimensions are in channel-last order.

When creating a NetworkRT object, you can specify the dimension format: CHW (default), NCHW, or ~NHWC. The latter is useful for TensorRT files that were exported from keras~

I've removed channels last format, since this is built in for intermediate layers. So this recent commit only specifies whether the TensorRT engine includes the batch dimension

nightduck commented 3 years ago

In retrospect, this PR is garbage, don't pull it. But the 1st commit is a good fix for issue #175, shall I remove the other two? Or just trash this PR and provide a patch instead?