Source code for sparseml.optim.analyzer

# Copyright (c) 2021 - present / Neuralmagic, Inc. 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.

"""
Code for describing layers / operators in ML framework neural networks.
"""

import json
from typing import Any, Dict, List, Tuple

from sparseml.utils import clean_path, create_parent_dirs


__all__ = ["AnalyzedLayerDesc"]


[docs]class AnalyzedLayerDesc(object): """ Description of an executed neural network layer. Contains information about the number of flops, shapes, params, etc. :param name: name of the layer :param type_: type of the layer :param params: number of parameters of the layer :param zeroed_params: number of parameters with values of zero :param prunable_params: number of parameters that could be pruned :param params_dims: dimensions of parameters :param prunable_params_dims: dimensions of prunable parameters :param execution_order: execution order of the layer/operation :param input_shape: shapes of input tensors :param output_shape: shapes of output tensors :param flops: Unused :param total_flops: total number of float operations """
[docs] @staticmethod def save_descs(descs: List, path: str): """ Save a list of AnalyzedLayerDesc to a json file :param descs: a list of descriptions to save :param path: the path to save the descriptions at """ path = clean_path(path) create_parent_dirs(path) save_obj = {"descriptions": [desc.dict() for desc in descs]} with open(path, "w") as file: json.dump(save_obj, file)
[docs] @staticmethod def load_descs(path: str) -> List: """ Load a list of AnalyzedLayerDesc from a json file :param path: the path to load the descriptions from :return: the loaded list of AnalyzedLayerDesc """ path = clean_path(path) with open(path, "r") as file: obj = json.load(file) descs = [] for desc_obj in obj["descriptions"]: desc_obj["type_"] = desc_obj["type"] del desc_obj["type"] del desc_obj["terminal"] del desc_obj["prunable"] descs.append(AnalyzedLayerDesc(**desc_obj)) return descs
[docs] @staticmethod def merge_descs(orig, descs: List): """ Merge a layer description with a list of others :param orig: original description :param descs: list of descriptions to merge with :return: a combined description """ merged = AnalyzedLayerDesc( name=orig.name, type_=orig.type_, params=orig.params, zeroed_params=orig.zeroed_params, prunable_params=orig.prunable_params, params_dims=orig.params_dims, prunable_params_dims=orig.prunable_params_dims, execution_order=orig.execution_order, input_shape=orig.input_shape, output_shape=orig.output_shape, flops=orig.flops, total_flops=orig.total_flops, stride=orig.stride, ) for desc in descs: merged.flops += desc.flops merged.total_flops += desc.total_flops merged.params += desc.params merged.prunable_params += desc.prunable_params merged.zeroed_params += desc.zeroed_params return merged
def __init__( self, name: str, type_: str, params: int = 0, zeroed_params: int = 0, prunable_params: int = 0, params_dims: Dict[str, Tuple[int, ...]] = None, prunable_params_dims: Dict[str, Tuple[int, ...]] = None, execution_order: int = -1, input_shape: Tuple[Tuple[int, ...], ...] = None, output_shape: Tuple[Tuple[int, ...], ...] = None, flops: int = 0, total_flops: int = 0, stride: Tuple[int, ...] = None, ): self.name = name self.type_ = type_ self.params = params self.prunable_params = prunable_params self.zeroed_params = zeroed_params self.params_dims = params_dims self.prunable_params_dims = prunable_params_dims self.execution_order = execution_order self.input_shape = input_shape self.output_shape = output_shape self.flops = flops self.total_flops = total_flops self.stride = stride def __repr__(self): return "AnalyzedLayerDesc({})".format(self.dict()) @property def terminal(self) -> bool: """ :return: True if this is a terminal op, ie it is doing compute and is not a container, False otherwise """ return self.params_dims is not None @property def prunable(self) -> bool: """ :return: True if the layer supports kernel sparsity (is prunable), False otherwise """ return self.prunable_params > 0
[docs] def dict(self) -> Dict[str, Any]: """ :return: A serializable dictionary representation of the current instance """ return { "name": self.name, "type": self.type_, "params": self.params, "zeroed_params": self.zeroed_params, "prunable_params": self.prunable_params, "params_dims": self.params_dims, "prunable_params_dims": self.prunable_params_dims, "execution_order": self.execution_order, "input_shape": self.input_shape, "output_shape": self.output_shape, "stride": self.stride, "flops": self.flops, "total_flops": self.total_flops, "terminal": self.terminal, "prunable": self.prunable, }