mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

yolo4tiny works on tensorRT 🐬 🐬 🐬 🐬

Browse Source 
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
Micaela Verucchi
2020-06-30 15:52:58 +02:00
parent 61aa24c6b7
commit fe2e4eae92
4 changed files with 27 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/NetworkRT.h
+1 -1
View File
@@ -28,7 +28,7 @@ using namespace nvinfer1;
#include "pluginsRT/ActivationMishRT.h"
#include "pluginsRT/ReorgRT.h"
#include "pluginsRT/RegionRT.h"
//#include "pluginsRT/RouteRT.h"
#include "pluginsRT/RouteRT.h"
#include "pluginsRT/ShortcutRT.h"
#include "pluginsRT/YoloRT.h"
#include "pluginsRT/UpsampleRT.h"
include/tkDNN/pluginsRT/RouteRT.h
+12 -5
View File
@@ -8,7 +8,9 @@ class RouteRT : public IPlugin {
*/
public:
RouteRT() {
RouteRT(int groups, int group_id) {
this->groups = groups;
this->group_id = group_id;
}
~RouteRT(){
@@ -22,7 +24,7 @@ public:
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
int out_c = 0;
for(int i=0; i<nbInputDims; i++) out_c += inputs[i].d[0];
return DimsCHW{out_c, inputs[0].d[1], inputs[0].d[2]};
return DimsCHW{out_c/groups, inputs[0].d[1], inputs[0].d[2]};
}
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
@@ -34,6 +36,7 @@ public:
}
h = inputDims[0].d[1];
w = inputDims[0].d[2];
c /= groups;
}
int initialize() override {
@@ -56,8 +59,9 @@ public:
for(int i=0; i<in; i++) {
dnnType *input = (dnnType*)reinterpret_cast<const dnnType*>(inputs[i]);
int in_dim = c_in[i]*h*w;
checkCuda( cudaMemcpyAsync(dstData + offset, input, in_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream) );
offset += in_dim;
int part_in_dim = in_dim / this->groups;
checkCuda( cudaMemcpyAsync(dstData + offset, input + this->group_id*part_in_dim, part_in_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream) );
offset += part_in_dim;
}
return 0;
@@ -65,11 +69,13 @@ public:
virtual size_t getSerializationSize() override {
return (4+MAX_INPUTS)*sizeof(int);
return (6+MAX_INPUTS)*sizeof(int);
}
virtual void serialize(void* buffer) override {
char *buf = reinterpret_cast<char*>(buffer);
tk::dnn::writeBUF(buf, groups);
tk::dnn::writeBUF(buf, group_id);
tk::dnn::writeBUF(buf, in);
for(int i=0; i<MAX_INPUTS; i++)
tk::dnn::writeBUF(buf, c_in[i]);
@@ -83,4 +89,5 @@ public:
int in;
int c_in[MAX_INPUTS];
int c, h, w;
int groups, group_id;
};