diff options
Diffstat (limited to 'src/graph/nodes')
| -rw-r--r-- | src/graph/nodes/ConvolutionLayerNode.cpp | 10 | ||||
| -rw-r--r-- | src/graph/nodes/SoftmaxLayerNode.cpp | 5 |
2 files changed, 12 insertions, 3 deletions
diff --git a/src/graph/nodes/ConvolutionLayerNode.cpp b/src/graph/nodes/ConvolutionLayerNode.cpp index 461728487f..eb0c6a1c1a 100644 --- a/src/graph/nodes/ConvolutionLayerNode.cpp +++ b/src/graph/nodes/ConvolutionLayerNode.cpp @@ -31,8 +31,8 @@ namespace arm_compute { namespace graph { -ConvolutionLayerNode::ConvolutionLayerNode(PadStrideInfo info, ConvolutionMethod method) - : _info(std::move(info)), _method(method) +ConvolutionLayerNode::ConvolutionLayerNode(PadStrideInfo info, ConvolutionMethod method, QuantizationInfo out_quant_info) + : _info(std::move(info)), _method(method), _out_quant_info(out_quant_info) { _input_edges.resize(3, EmptyEdgeID); _outputs.resize(1, NullTensorID); @@ -90,6 +90,12 @@ TensorDescriptor ConvolutionLayerNode::configure_output(size_t idx) const TensorDescriptor output_info = src->desc(); TensorShape output_shape = compute_output_shape(src->desc().shape, weights->desc().shape, _info); output_info.shape = output_shape; + + if(!_out_quant_info.empty()) + { + output_info.quant_info = _out_quant_info; + } + return output_info; } diff --git a/src/graph/nodes/SoftmaxLayerNode.cpp b/src/graph/nodes/SoftmaxLayerNode.cpp index 4c21ac6ad0..b6241e6654 100644 --- a/src/graph/nodes/SoftmaxLayerNode.cpp +++ b/src/graph/nodes/SoftmaxLayerNode.cpp @@ -63,7 +63,10 @@ TensorDescriptor SoftmaxLayerNode::configure_output(size_t idx) const const Tensor *src = input(0); ARM_COMPUTE_ERROR_ON(src == nullptr); - return src->desc(); + TensorDescriptor out_desc = src->desc(); + out_desc.quant_info = QuantizationInfo(1.f / 256.f, 0); + + return out_desc; } Status SoftmaxLayerNode::validate() |