tkDNN
tkDNN copied to clipboard
ceccocats
Additional support for custom network architectures
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)
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.
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
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?