105 lines
3.9 KiB
C++
105 lines
3.9 KiB
C++
/*
|
|
* Created by Marcos Luciano
|
|
* https://www.github.com/marcoslucianops
|
|
*/
|
|
|
|
#include "activation_layer.h"
|
|
|
|
nvinfer1::ILayer* activationLayer(
|
|
int layerIdx,
|
|
std::string activation,
|
|
nvinfer1::ILayer* output,
|
|
nvinfer1::ITensor* input,
|
|
nvinfer1::INetworkDefinition* network)
|
|
{
|
|
if (activation == "linear") {
|
|
// Pass
|
|
}
|
|
else if (activation == "relu")
|
|
{
|
|
nvinfer1::IActivationLayer* relu = network->addActivation(
|
|
*input, nvinfer1::ActivationType::kRELU);
|
|
assert(relu != nullptr);
|
|
std::string reluLayerName = "relu_" + std::to_string(layerIdx);
|
|
relu->setName(reluLayerName.c_str());
|
|
output = relu;
|
|
}
|
|
else if (activation == "sigmoid" || activation == "logistic")
|
|
{
|
|
nvinfer1::IActivationLayer* sigmoid = network->addActivation(
|
|
*input, nvinfer1::ActivationType::kSIGMOID);
|
|
assert(sigmoid != nullptr);
|
|
std::string sigmoidLayerName = "sigmoid_" + std::to_string(layerIdx);
|
|
sigmoid->setName(sigmoidLayerName.c_str());
|
|
output = sigmoid;
|
|
}
|
|
else if (activation == "tanh")
|
|
{
|
|
nvinfer1::IActivationLayer* tanh = network->addActivation(
|
|
*input, nvinfer1::ActivationType::kTANH);
|
|
assert(tanh != nullptr);
|
|
std::string tanhLayerName = "tanh_" + std::to_string(layerIdx);
|
|
tanh->setName(tanhLayerName.c_str());
|
|
output = tanh;
|
|
}
|
|
else if (activation == "leaky")
|
|
{
|
|
nvinfer1::IActivationLayer* leaky = network->addActivation(
|
|
*input, nvinfer1::ActivationType::kLEAKY_RELU);
|
|
leaky->setAlpha(0.1);
|
|
assert(leaky != nullptr);
|
|
std::string leakyLayerName = "leaky_" + std::to_string(layerIdx);
|
|
leaky->setName(leakyLayerName.c_str());
|
|
output = leaky;
|
|
}
|
|
else if (activation == "softplus")
|
|
{
|
|
nvinfer1::IActivationLayer* softplus = network->addActivation(
|
|
*input, nvinfer1::ActivationType::kSOFTPLUS);
|
|
assert(softplus != nullptr);
|
|
std::string softplusLayerName = "softplus_" + std::to_string(layerIdx);
|
|
softplus->setName(softplusLayerName.c_str());
|
|
output = softplus;
|
|
}
|
|
else if (activation == "mish")
|
|
{
|
|
nvinfer1::IActivationLayer* softplus = network->addActivation(
|
|
*input, nvinfer1::ActivationType::kSOFTPLUS);
|
|
assert(softplus != nullptr);
|
|
std::string softplusLayerName = "softplus_" + std::to_string(layerIdx);
|
|
softplus->setName(softplusLayerName.c_str());
|
|
nvinfer1::IActivationLayer* tanh = network->addActivation(
|
|
*softplus->getOutput(0), nvinfer1::ActivationType::kTANH);
|
|
assert(tanh != nullptr);
|
|
std::string tanhLayerName = "tanh_" + std::to_string(layerIdx);
|
|
tanh->setName(tanhLayerName.c_str());
|
|
nvinfer1::IElementWiseLayer* mish = network->addElementWise(
|
|
*tanh->getOutput(0), *input,
|
|
nvinfer1::ElementWiseOperation::kPROD);
|
|
assert(mish != nullptr);
|
|
std::string mishLayerName = "mish_" + std::to_string(layerIdx);
|
|
mish->setName(mishLayerName.c_str());
|
|
output = mish;
|
|
}
|
|
else if (activation == "silu" || activation == "swish")
|
|
{
|
|
nvinfer1::IActivationLayer* sigmoid = network->addActivation(
|
|
*input, nvinfer1::ActivationType::kSIGMOID);
|
|
assert(sigmoid != nullptr);
|
|
std::string sigmoidLayerName = "sigmoid_" + std::to_string(layerIdx);
|
|
sigmoid->setName(sigmoidLayerName.c_str());
|
|
nvinfer1::IElementWiseLayer* silu = network->addElementWise(
|
|
*sigmoid->getOutput(0), *input,
|
|
nvinfer1::ElementWiseOperation::kPROD);
|
|
assert(silu != nullptr);
|
|
std::string siluLayerName = "silu_" + std::to_string(layerIdx);
|
|
silu->setName(siluLayerName.c_str());
|
|
output = silu;
|
|
}
|
|
else {
|
|
std::cerr << "Activation not supported: " << activation << std::endl;
|
|
std::abort();
|
|
}
|
|
return output;
|
|
}
|