diff options
Diffstat (limited to 'src/armnn/NetworkUtils.cpp')
-rw-r--r-- | src/armnn/NetworkUtils.cpp | 29 |
1 files changed, 29 insertions, 0 deletions
diff --git a/src/armnn/NetworkUtils.cpp b/src/armnn/NetworkUtils.cpp index 1e3add6301..9a4ce87b59 100644 --- a/src/armnn/NetworkUtils.cpp +++ b/src/armnn/NetworkUtils.cpp @@ -74,4 +74,33 @@ std::vector<ConvertFp32ToFp16Layer*> InsertConvertFp32ToFp16LayersAfter(Graph& g return convertLayers; } + +std::vector<DebugLayer*> InsertDebugLayerAfter(Graph& graph, Layer& layer) +{ + std::vector<DebugLayer*> debugLayers; + debugLayers.reserve(layer.GetNumOutputSlots()); + + // Change outputs to DataType::Float16 + for (auto&& outputSlot = layer.BeginOutputSlots(); outputSlot != layer.EndOutputSlots(); ++outputSlot) + { + // Insert debug layer after the layer + const std::string name = + std::string("DebugLayerAfter") + layer.GetName(); + + const DebugDescriptor descriptor; + + DebugLayer* debugLayer = + graph.InsertNewLayer<DebugLayer>(*outputSlot, descriptor, name.c_str()); + + // Sets output tensor info for the debug layer. + TensorInfo debugInfo = debugLayer->GetInputSlot(0).GetConnectedOutputSlot()->GetTensorInfo(); + + debugLayer->GetOutputSlot().SetTensorInfo(debugInfo); + + debugLayers.emplace_back(debugLayer); + } + + return debugLayers; +} + } // namespace armnn |