diff options
Diffstat (limited to 'src/armnn/QuantizerVisitor.cpp')
| -rw-r--r-- | src/armnn/QuantizerVisitor.cpp | 20 |
1 files changed, 12 insertions, 8 deletions
diff --git a/src/armnn/QuantizerVisitor.cpp b/src/armnn/QuantizerVisitor.cpp index 4b80b02e34..51818ebddd 100644 --- a/src/armnn/QuantizerVisitor.cpp +++ b/src/armnn/QuantizerVisitor.cpp @@ -115,9 +115,7 @@ void QuantizerVisitor::RecordLayer(const IConnectableLayer* srcLayer, IConnectab void QuantizerVisitor::VisitAbsLayer(const IConnectableLayer* layer, const char* name) { - IConnectableLayer* newLayer = m_QuantizedNetwork->AddAbsLayer(name); - RecordLayer(layer, newLayer); - SetQuantizedInputConnections(layer, newLayer); + VisitElementwiseUnaryLayer(layer, ElementwiseUnaryDescriptor(UnaryOperation::Abs), name); } void QuantizerVisitor::VisitActivationLayer(const IConnectableLayer* layer, @@ -275,6 +273,15 @@ void QuantizerVisitor::VisitDepthwiseConvolution2dLayer(const IConnectableLayer* SetQuantizedInputConnections(layer, newLayer); } +void QuantizerVisitor::VisitElementwiseUnaryLayer(const IConnectableLayer* layer, + const ElementwiseUnaryDescriptor& elementwiseUnaryDescriptor, + const char* name) +{ + IConnectableLayer* newLayer = m_QuantizedNetwork->AddElementwiseUnaryLayer(elementwiseUnaryDescriptor, name); + RecordLayer(layer, newLayer); + SetQuantizedInputConnections(layer, newLayer); +} + void QuantizerVisitor::VisitFullyConnectedLayer(const IConnectableLayer *layer, const FullyConnectedDescriptor& desc, const ConstTensor& weights, @@ -450,12 +457,9 @@ void QuantizerVisitor::VisitResizeLayer(const IConnectableLayer* layer, SetQuantizedInputConnections(layer, newLayer); } -void QuantizerVisitor::VisitRsqrtLayer(const IConnectableLayer* layer, - const char* name) +void QuantizerVisitor::VisitRsqrtLayer(const IConnectableLayer* layer, const char* name) { - IConnectableLayer* newLayer = m_QuantizedNetwork->AddRsqrtLayer(name); - RecordLayer(layer, newLayer); - SetQuantizedInputConnections(layer, newLayer); + VisitElementwiseUnaryLayer(layer, ElementwiseUnaryDescriptor(UnaryOperation::Rsqrt), name); } void QuantizerVisitor::VisitSliceLayer(const IConnectableLayer* layer, |