diff --git a/include/Layer.h b/include/Layer.h index 06b28f6..ab14a43 100644 --- a/include/Layer.h +++ b/include/Layer.h @@ -97,6 +97,8 @@ typedef enum { class Activation : public Layer { public: + int act_mode; + Activation(Network *net, int act_mode); virtual ~Activation(); virtual layerType_t getLayerType() { return LAYER_ACTIVATION; }; @@ -104,7 +106,6 @@ public: virtual value_type* infer(dataDim_t &dim, value_type* srcData); protected: - int act_mode; cudnnActivationDescriptor_t activDesc; }; diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 23ed901..3551cfc 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -4,6 +4,7 @@ #include "NetworkRT.h" using namespace nvinfer1; +#include "pluginsRT/ActivationLeakyRT.cpp" // Logger for info/warning/errors class Logger : public ILogger @@ -183,14 +184,21 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Pooling *l) { ITensor* NetworkRT::convert_layer(ITensor *input, Activation *l) { std::cout<<"convert Activation\n"; + if(l->act_mode == ACTIVATION_LEAKY) { + std::cout<<"New plugin LEAKY\n"; + IPlugin *plugin = new ActivationLeakyRT(); + IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); + checkNULL(lRT); + return lRT->getOutput(0); + } + IActivationLayer *lRT = networkRT->addActivation(*input, ActivationType::kRELU); checkNULL(lRT); - return lRT->getOutput(0); } ITensor* NetworkRT::convert_layer(ITensor *input, Softmax *l) { - std::cout<<"convert Activation\n"; + std::cout<<"convert softmax\n"; ISoftMaxLayer *lRT = networkRT->addSoftMax(*input); checkNULL(lRT); diff --git a/src/pluginsRT/ActivationLeakyRT.cpp b/src/pluginsRT/ActivationLeakyRT.cpp new file mode 100644 index 0000000..355d8ea --- /dev/null +++ b/src/pluginsRT/ActivationLeakyRT.cpp @@ -0,0 +1,58 @@ +#include +#include "kernels.h" + +class ActivationLeakyRT : public IPlugin { + +public: + ActivationLeakyRT() { + + + } + + ~ActivationLeakyRT(){ + + } + + int getNbOutputs() const override { + return 1; + } + + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override { + return inputs[0]; + } + + void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override { + size = 1; + for(int i=0; i(inputs[0]), + reinterpret_cast(outputs[0]), size); + return 0; + } + + + virtual size_t getSerializationSize() override { + return 0; + } + + virtual void serialize(void* buffer) override { + } + + int size; +}; diff --git a/tests/mnist/test_mnist.cpp b/tests/mnist/test_mnist.cpp index f85cfcd..b5dbead 100644 --- a/tests/mnist/test_mnist.cpp +++ b/tests/mnist/test_mnist.cpp @@ -18,7 +18,7 @@ int main() { tkDNN::Conv2d l2(&net, 50, 5, 5, 1, 1, 0, 0, c1_bin); tkDNN::Pooling l3(&net, 2, 2, 2, 2, tkDNN::POOLING_MAX); tkDNN::Dense l4(&net, 500, d2_bin); - tkDNN::Activation l5(&net, CUDNN_ACTIVATION_RELU); + tkDNN::Activation l5(&net, tkDNN::ACTIVATION_LEAKY); tkDNN::Dense l6(&net, 10, d3_bin); tkDNN::Softmax l7(&net); diff --git a/tests/yolo/yolo.cpp b/tests/yolo/yolo.cpp index 9e4eadb..04d15dc 100644 --- a/tests/yolo/yolo.cpp +++ b/tests/yolo/yolo.cpp @@ -83,7 +83,7 @@ int main() { tkDNN::Activation a23(&net, tkDNN::ACTIVATION_LEAKY); tkDNN::Conv2d c24(&net, 1024, 3, 3, 1, 1, 1, 1, c24_bin, true); tkDNN::Activation a24(&net, tkDNN::ACTIVATION_LEAKY); - +/* tkDNN::Layer *m25_layers[1] = { &a16 }; tkDNN::Route m25(&net, m25_layers, 1); tkDNN::Conv2d c26(&net, 64, 1, 1, 1, 1, 0, 0, c26_bin, true); @@ -98,7 +98,7 @@ int main() { tkDNN::Conv2d c30(&net, 425, 1, 1, 1, 1, 0, 0, c30_bin, false); tkDNN::Region g31(&net, 80, 4, 5, 0.6f); - +*/ // Load input value_type *data; value_type *input_h; @@ -108,16 +108,16 @@ int main() { value_type *out_data, *out_data2; + tkDNN::dataDim_t dim1 = dim; std::cout<<"CUDNN inference:\n"; { - dim.print(); //print initial dimension + dim1.print(); //print initial dimension TIMER_START - out_data = net.infer(dim, data); + out_data = net.infer(dim1, data); TIMER_STOP - dim.print(); + dim1.print(); } - tkDNN::dataDim_t dim2(1, 3, 608, 608, 1); - + tkDNN::dataDim_t dim2 = dim; std::cout<<"TENSORRT inference:\n"; { dim2.print(); TIMER_START @@ -127,7 +127,7 @@ int main() { } std::cout<<"\n======= CHECK RESULT =======\n"; - std::cout<<"Diffs: "<