Spaces:
Running
Running
| # Copyright 2016 The TensorFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Graph post-processing logic. | |
| Used by both TensorBoard and mldash. | |
| """ | |
| def prepare_graph_for_ui( | |
| graph, limit_attr_size=1024, large_attrs_key="_too_large_attrs" | |
| ): | |
| """Prepares (modifies in-place) the graph to be served to the front-end. | |
| For now, it supports filtering out attributes that are | |
| too large to be shown in the graph UI. | |
| Args: | |
| graph: The GraphDef proto message. | |
| limit_attr_size: Maximum allowed size in bytes, before the attribute | |
| is considered large. Default is 1024 (1KB). Must be > 0 or None. | |
| If None, there will be no filtering. | |
| large_attrs_key: The attribute key that will be used for storing attributes | |
| that are too large. Default is '_too_large_attrs'. Must be != None if | |
| `limit_attr_size` is != None. | |
| Raises: | |
| ValueError: If `large_attrs_key is None` while `limit_attr_size != None`. | |
| ValueError: If `limit_attr_size` is defined, but <= 0. | |
| """ | |
| # TODO(@davidsoergel): detect whether a graph has been filtered already | |
| # (to a limit_attr_size <= what is requested here). If it is already | |
| # filtered, return immediately. | |
| # Check input for validity. | |
| if limit_attr_size is not None: | |
| if large_attrs_key is None: | |
| raise ValueError( | |
| "large_attrs_key must be != None when limit_attr_size" | |
| "!= None." | |
| ) | |
| if limit_attr_size <= 0: | |
| raise ValueError( | |
| "limit_attr_size must be > 0, but is %d" % limit_attr_size | |
| ) | |
| # Filter only if a limit size is defined. | |
| if limit_attr_size is not None: | |
| for node in graph.node: | |
| # Go through all the attributes and filter out ones bigger than the | |
| # limit. | |
| keys = list(node.attr.keys()) | |
| for key in keys: | |
| size = node.attr[key].ByteSize() | |
| if size > limit_attr_size or size < 0: | |
| del node.attr[key] | |
| # Add the attribute key to the list of "too large" attributes. | |
| # This is used in the info card in the graph UI to show the user | |
| # that some attributes are too large to be shown. | |
| node.attr[large_attrs_key].list.s.append( | |
| key.encode("utf-8") | |
| ) | |