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

Shelfnet works on cuDNN. To test everything else

Browse Source 
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
Micaela Verucchi
2020-06-22 18:11:19 +02:00
parent 9f1e30eaa9
commit 6fd261f628
11 changed files with 290 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/pluginsRT/ShortcutRT.h
+7 -5
View File
@@ -4,10 +4,11 @@
class ShortcutRT : public IPlugin {
public:
ShortcutRT(tk::dnn::dataDim_t bdim) {
ShortcutRT(tk::dnn::dataDim_t bdim, bool mul) {
this->bc = bdim.c;
this->bh = bdim.h;
this->bw = bdim.w;
this->mul = mul;
}
~ShortcutRT(){
@@ -47,15 +48,14 @@ public:
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
for(int b=0; b < batchSize; ++b)
shortcutForward(srcDataBack + b*bc*bh*bw, dstData + b*c*h*w, 1, c, h, w, 1, 1, bc, bh, bw, 1, stream);
shortcutForward(srcDataBack, dstData, batchSize, c, h, w, 1, batchSize, bc, bh, bw, 1, mul, stream);
return 0;
}
virtual size_t getSerializationSize() override {
return 6*sizeof(int);
return 6*sizeof(int) + sizeof(bool);
}
virtual void serialize(void* buffer) override {
@@ -63,12 +63,14 @@ public:
tk::dnn::writeBUF(buf, bc);
tk::dnn::writeBUF(buf, bh);
tk::dnn::writeBUF(buf, bw);
tk::dnn::writeBUF(buf, mul);
tk::dnn::writeBUF(buf, c);
tk::dnn::writeBUF(buf, h);
tk::dnn::writeBUF(buf, w);
}
int c, h, w;
int bc, bh, bw;
bool mul;
};