Big update
This commit is contained in:
@@ -8,7 +8,7 @@
|
||||
#include <cassert>
|
||||
|
||||
nvinfer1::ITensor*
|
||||
shortcutLayer(int layerIdx, std::string mode, std::string activation, std::string inputVol, std::string shortcutVol,
|
||||
shortcutLayer(int layerIdx, std::string activation, std::string inputVol, std::string shortcutVol,
|
||||
std::map<std::string, std::string>& block, nvinfer1::ITensor* input, nvinfer1::ITensor* shortcutInput,
|
||||
nvinfer1::INetworkDefinition* network)
|
||||
{
|
||||
@@ -16,12 +16,7 @@ shortcutLayer(int layerIdx, std::string mode, std::string activation, std::strin
|
||||
|
||||
assert(block.at("type") == "shortcut");
|
||||
|
||||
nvinfer1::ElementWiseOperation operation = nvinfer1::ElementWiseOperation::kSUM;
|
||||
|
||||
if (mode == "mul")
|
||||
operation = nvinfer1::ElementWiseOperation::kPROD;
|
||||
|
||||
if (mode == "add" && inputVol != shortcutVol) {
|
||||
if (inputVol != shortcutVol) {
|
||||
nvinfer1::ISliceLayer* slice = network->addSlice(*shortcutInput, nvinfer1::Dims{3, {0, 0, 0}}, input->getDimensions(),
|
||||
nvinfer1::Dims{3, {1, 1, 1}});
|
||||
assert(slice != nullptr);
|
||||
@@ -32,7 +27,7 @@ shortcutLayer(int layerIdx, std::string mode, std::string activation, std::strin
|
||||
else
|
||||
output = shortcutInput;
|
||||
|
||||
nvinfer1::IElementWiseLayer* shortcut = network->addElementWise(*input, *output, operation);
|
||||
nvinfer1::IElementWiseLayer* shortcut = network->addElementWise(*input, *output, nvinfer1::ElementWiseOperation::kSUM);
|
||||
assert(shortcut != nullptr);
|
||||
std::string shortcutLayerName = "shortcut_" + std::to_string(layerIdx);
|
||||
shortcut->setName(shortcutLayerName.c_str());
|
||||
|
||||
Reference in New Issue
Block a user