Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 1 | # Copyright (C) 2021 Arm Limited or its affiliates. All rights reserved. |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 2 | # |
| 3 | # SPDX-License-Identifier: Apache-2.0 |
| 4 | # |
| 5 | # Licensed under the Apache License, Version 2.0 (the License); you may |
| 6 | # not use this file except in compliance with the License. |
| 7 | # You may obtain a copy of the License at |
| 8 | # |
| 9 | # www.apache.org/licenses/LICENSE-2.0 |
| 10 | # |
| 11 | # Unless required by applicable law or agreed to in writing, software |
| 12 | # distributed under the License is distributed on an AS IS BASIS, WITHOUT |
| 13 | # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 16 | # Description: |
Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 17 | # Early optimisation of the network graph, using the rewrite_graph module to do the traversal of the graph. |
Diego Russo | ea6111a | 2020-04-14 18:41:58 +0100 | [diff] [blame] | 18 | from . import rewrite_graph |
Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 19 | from .graph_optimiser_util import check_format_restrictions |
Jonas Ohlsson | 0957e3e | 2021-09-01 15:57:21 +0200 | [diff] [blame] | 20 | from .graph_optimiser_util import check_memory_only_removed |
Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 21 | from .graph_optimiser_util import record_optimised |
| 22 | from .nn_graph import NetworkType |
| 23 | from .tflite_graph_optimiser import tflite_optimise_graph |
| 24 | from .tosa_graph_optimiser import tosa_optimise_graph |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 25 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 26 | |
Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 27 | def optimise_graph(nng, arch, network_type, verbose_graph=False): |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 28 | if verbose_graph: |
Fredrik Svedberg | c875aa6 | 2021-05-06 09:53:31 +0200 | [diff] [blame] | 29 | nng.print_graph("Before Graph Optimization") |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 30 | |
Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 31 | if network_type == NetworkType.TFLite: |
| 32 | # TensorFlow Lite graph optimization |
| 33 | nng = tflite_optimise_graph(nng, arch) |
| 34 | else: |
| 35 | # TOSA graph optimization |
| 36 | nng = tosa_optimise_graph(nng, arch) |
Patrik Gustavsson | ee99bb1 | 2021-04-08 09:04:00 +0200 | [diff] [blame] | 37 | |
Patrik Gustavsson | 3a26920 | 2021-01-21 08:28:55 +0100 | [diff] [blame] | 38 | # Post-optimisation operator debug tracing, and checking that no undesired reshapes are left in the graph |
Tim Hall | e6ccd87 | 2020-11-09 16:46:37 +0000 | [diff] [blame] | 39 | for sg in nng.subgraphs: |
Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 40 | rewrite_graph.visit_graph_post_order( |
Jonas Ohlsson | 0957e3e | 2021-09-01 15:57:21 +0200 | [diff] [blame] | 41 | sg.output_tensors, arch, [check_format_restrictions], [check_memory_only_removed, record_optimised] |
Patrik Gustavsson | 8f1f9aa | 2021-06-28 07:41:58 +0200 | [diff] [blame] | 42 | ) |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 43 | |
| 44 | if verbose_graph: |
Fredrik Svedberg | c875aa6 | 2021-05-06 09:53:31 +0200 | [diff] [blame] | 45 | nng.print_graph("After Graph Optimization") |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 46 | return nng |