Add swish

include/tkDNN/Layer.h

@@ -19,6 +19,7 @@ enum layerType_t {
     LAYER_ACTIVATION_CRELU,
     LAYER_ACTIVATION_LEAKY,
     LAYER_ACTIVATION_MISH,
+    LAYER_ACTIVATION_SWISH,
     LAYER_ACTIVATION_LOGISTIC,
     LAYER_FLATTEN,
     LAYER_RESHAPE,
@@ -75,6 +76,7 @@ public:
             case LAYER_ACTIVATION_CRELU:    return "ActivationCReLU";
             case LAYER_ACTIVATION_LEAKY:    return "ActivationLeaky";
             case LAYER_ACTIVATION_MISH:     return "ActivationMish";
+            case LAYER_ACTIVATION_SWISH:    return "ActivationSwish";
             case LAYER_ACTIVATION_LOGISTIC: return "ActivationLogistic";
             case LAYER_FLATTEN:             return "Flatten";
             case LAYER_RESHAPE:             return "Reshape";
@@ -223,7 +225,8 @@ typedef enum {
     ACTIVATION_ELU     = 100,
     ACTIVATION_LEAKY   = 101,
     ACTIVATION_MISH   = 102,
-    ACTIVATION_LOGISTIC   = 103
+    ACTIVATION_LOGISTIC   = 103,
+    ACTIVATION_SWISH = 104
 } tkdnnActivationMode_t;
 
 /**
@@ -245,6 +248,8 @@ public:
             return LAYER_ACTIVATION_LEAKY;
         else if (act_mode == ACTIVATION_MISH)
             return LAYER_ACTIVATION_MISH;
+        else if (act_mode == ACTIVATION_SWISH)
+            return LAYER_ACTIVATION_SWISH;
         else if (act_mode == ACTIVATION_LOGISTIC)
             return LAYER_ACTIVATION_LOGISTIC;
         else

include/tkDNN/NetworkRT.h

@@ -26,6 +26,47 @@
 #include <pluginsRT/ConstantPaddingRT.h>
 #include <pluginsRT/ReflectionPadding.h>
 
+namespace tk { namespace dnn {
+
+template<typename T> void writeBUF(char*& buffer, const T& val)
+{
+    *reinterpret_cast<T*>(buffer) = val;
+    buffer += sizeof(T);
+}
+
+template<typename T> T readBUF(const char*& buffer)
+{
+    T val = *reinterpret_cast<const T*>(buffer);
+    buffer += sizeof(T);
+    return val;
+}
+
+using namespace nvinfer1;
+#include "pluginsRT/ActivationLeakyRT.h"
+#include "pluginsRT/ActivationLogisticRT.h"
+#include "pluginsRT/ActivationReLUCeilingRT.h"
+#include "pluginsRT/ActivationMishRT.h"
+#include "pluginsRT/ActivationSwishRT.h"
+#include "pluginsRT/ReorgRT.h"
+#include "pluginsRT/RegionRT.h"
+#include "pluginsRT/RouteRT.h"
+#include "pluginsRT/ShortcutRT.h"
+#include "pluginsRT/YoloRT.h"
+#include "pluginsRT/UpsampleRT.h"
+#include "pluginsRT/ResizeLayerRT.h"
+#include "pluginsRT/DeformableConvRT.h"
+#include "pluginsRT/FlattenConcatRT.h"
+#include "pluginsRT/ReshapeRT.h"
+#include "pluginsRT/MaxPoolingFixedSizeRT.h"
+
+class PluginFactory : IPluginFactory
+{
+public:
+    YoloRT *yolos[16];
+    int n_yolos;
+
+	virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength);
+};
 
 
 namespace tk { namespace dnn {

include/tkDNN/kernels.h

@@ -9,6 +9,7 @@ void activationReLUCeilingForward(dnnType *srcData, dnnType *dstData, int size,
 void activationLOGISTICForward(dnnType *srcData, dnnType *dstData, int size, cudaStream_t stream = cudaStream_t(0));
 void activationSIGMOIDForward(dnnType *srcData, dnnType *dstData, int size, cudaStream_t stream = cudaStream_t(0));
 void activationMishForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream= cudaStream_t(0));
+void activationSwishForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream= cudaStream_t(0));
 
 void fill(dnnType *data, int size, dnnType val, cudaStream_t stream = cudaStream_t(0));
 

include/tkDNN/pluginsRT/ActivationSwishRT.h

@@ -0,0 +1,61 @@
+#include<cassert>
+#include "../kernels.h"
+
+class ActivationSwishRT : public IPlugin {
+
+public:
+	ActivationSwishRT() {
+
+
+	}
+
+	~ActivationSwishRT(){
+
+	}
+
+	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<outputDims[0].nbDims; i++)
+			size *= outputDims[0].d[i];
+	}
+
+	int initialize() override {
+
+		return 0;
+	}
+
+	virtual void terminate() override {
+	}
+
+	virtual size_t getWorkspaceSize(int maxBatchSize) const override {
+		return 0;
+	}
+
+	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
+
+		activationSwishForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]), 
+											reinterpret_cast<dnnType*>(outputs[0]), batchSize*size, stream);
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 1*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer),*a=buf;
+		tk::dnn::writeBUF(buf, size);
+		assert(buf == a + getSerializationSize());
+	}
+
+	int size;
+};

src/Activation.cpp

@@ -52,6 +52,10 @@ dnnType* Activation::infer(dataDim_t &dim, dnnType* srcData) {
     else if(act_mode == ACTIVATION_MISH) {
         activationMishForward(srcData, dstData, dim.tot());
 
+    }
+    else if(act_mode == ACTIVATION_SWISH) {
+        activationSwishForward(srcData, dstData, dim.tot());
+
     }
     else if(act_mode == ACTIVATION_LOGISTIC) {
         activationLOGISTICForward(srcData, dstData, dim.tot());

src/DarknetParser.cpp

@@ -197,6 +197,7 @@ namespace tk { namespace dnn {
             if(f.activation == "relu") act = tkdnnActivationMode_t(CUDNN_ACTIVATION_RELU);
             else if(f.activation == "leaky") act = tk::dnn::ACTIVATION_LEAKY;
             else if(f.activation == "mish") act = tk::dnn::ACTIVATION_MISH;
+            else if(f.activation == "swish") act = tk::dnn::ACTIVATION_SWISH;
             else if(f.activation == "logistic") act = tk::dnn::ACTIVATION_LOGISTIC;
             else { FatalError("activation not supported: " + f.activation); }
             netLayers[netLayers.size()-1] = new tk::dnn::Activation(net, act);

src/NetworkRT.cpp

@@ -254,7 +254,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
         return convert_layer(input, (Conv2d*) l);
     if(type == LAYER_POOLING)
         return convert_layer(input, (Pooling*) l);
-    if(type == LAYER_ACTIVATION || type == LAYER_ACTIVATION_CRELU || type == LAYER_ACTIVATION_LEAKY || type == LAYER_ACTIVATION_MISH || type == LAYER_ACTIVATION_LOGISTIC)
+    if(type == LAYER_ACTIVATION || type == LAYER_ACTIVATION_CRELU || type == LAYER_ACTIVATION_LEAKY || type == LAYER_ACTIVATION_MISH || type == LAYER_ACTIVATION_SWISH || type == LAYER_ACTIVATION_LOGISTIC)
         return convert_layer(input, (Activation*) l);
     if(type == LAYER_SOFTMAX)
         return convert_layer(input, (Softmax*) l);
@@ -648,6 +648,15 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {
     }
     else if(l->act_mode == CUDNN_ACTIVATION_ELU || l->act_mode == ACTIVATION_ELU){
         IActivationLayer *lRT = networkRT->addActivation(*input,ActivationType::kELU);
+    else if(l->act_mode == ACTIVATION_SWISH) {
+        IPlugin *plugin = new ActivationSwishRT();
+        IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
+        checkNULL(lRT);
+        return lRT;
+    }
+    else if(l->act_mode == ACTIVATION_LOGISTIC) {
+        IPlugin *plugin = new ActivationLogisticRT();
+        IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
         checkNULL(lRT);
         return lRT;
     }
@@ -1030,6 +1039,36 @@ void NetworkRT::destroy() {
     if(builderActive) {
         delete engineRT;
         delete builderRT;
+
+
+IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData, size_t serialLength) {
+    const char * buf = reinterpret_cast<const char*>(serialData),*bufCheck = buf;
+
+    std::string name(layerName);
+    //std::cout<<name<<std::endl;
+
+    if(name.find("ActivationLeaky") == 0) {
+        ActivationLeakyRT *a = new ActivationLeakyRT(readBUF<float>(buf));
+        a->size = readBUF<int>(buf);
+        assert(buf == bufCheck + serialLength);
+        return a;
+    }
+    if(name.find("ActivationMish") == 0) {
+        ActivationMishRT *a = new ActivationMishRT();
+        a->size = readBUF<int>(buf);
+        assert(buf == bufCheck + serialLength);
+        return a;
+    }
+    if(name.find("ActivationSwish") == 0) {
+        ActivationSwishRT *a = new ActivationSwishRT();
+        a->size = readBUF<int>(buf);
+        assert(buf == bufCheck + serialLength);
+        return a;
+    }
+    if(name.find("ActivationLogistic") == 0) {
+        ActivationLogisticRT *a = new ActivationLogisticRT();
+        a->size = readBUF<int>(buf);
+        return a;
     }
 }
 #elif NV_TENSORRT_MAJOR <=7

src/kernels/activation_swish.cu

@@ -0,0 +1,23 @@
+#include "kernels.h"
+
+__global__
+void activation_swish(dnnType *input, dnnType *output, int size) {
+
+    int i = blockDim.x*blockIdx.x + threadIdx.x;
+
+    if(i<size) {    
+        output[i] =  input[i] * 1.0f/(1.0f + exp(-input[i]));;
+    }
+ }
+
+
+/**
+    Swish activation function
+*/
+void activationSwishForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream)
+{
+    int blocks = (size+255)/256;
+    int threads = 256;
+    
+    activation_swish<<<blocks, threads, 0, stream>>>(srcData, dstData, size);
+}