This repository has been archived on 2026-02-22. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
tkDNN/src/Shortcut.cpp
T
2020-06-23 12:50:24 +02:00

38 lines
1007 B
C++

#include <iostream>
#include "Layer.h"
#include "kernels.h"
namespace tk { namespace dnn {
Shortcut::Shortcut(Network *net, Layer *backLayer, bool mul) : Layer(net) {
this->backLayer = backLayer;
this->mul = mul;
checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
if( ( backLayer->output_dim.c != input_dim.c && mul ) ||
(( backLayer->output_dim.w != input_dim.w || backLayer->output_dim.h != input_dim.h ) && !mul ) )
FatalError("Shortcut dim missmatch");
}
Shortcut::~Shortcut() {
checkCuda( cudaFree(dstData) );
}
dnnType* Shortcut::infer(dataDim_t &dim, dnnType* srcData) {
dataDim_t bdim = this->backLayer->output_dim;
checkCuda(cudaMemcpy(dstData, srcData, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice));
shortcutForward(this->backLayer->dstData, dstData, dim.n, dim.c, dim.h, dim.w, 1, bdim.n, bdim.c, bdim.h, bdim.w, 1, mul);
//update data dimensions
dim = output_dim;
return dstData;
}
}}