# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- import logging from typing import Optional from fusion_attention_unet import FusionAttentionUnet from fusion_bias_add import FusionBiasAdd from fusion_biassplitgelu import FusionBiasSplitGelu from fusion_group_norm import FusionGroupNorm from fusion_nhwc_conv import FusionNhwcConv from fusion_options import FusionOptions from fusion_skip_group_norm import FusionSkipGroupNorm from fusion_transpose import FusionInsertTranspose, FusionTranspose from import_utils import is_installed from onnx import ModelProto from onnx_model import OnnxModel from onnx_model_bert import BertOnnxModel logger = logging.getLogger(__name__) class UnetOnnxModel(BertOnnxModel): def __init__(self, model: ModelProto, num_heads: int = 0, hidden_size: int = 0): """Initialize UNet ONNX Model. Args: model (ModelProto): the ONNX model num_heads (int, optional): number of attention heads. Defaults to 0 (detect the parameter automatically). hidden_size (int, optional): hidden dimension. Defaults to 0 (detect the parameter automatically). """ assert (num_heads == 0 and hidden_size == 0) or (num_heads > 0 and hidden_size % num_heads == 0) super().__init__(model, num_heads=num_heads, hidden_size=hidden_size) def preprocess(self): self.remove_useless_div() def postprocess(self): self.prune_graph() self.remove_unused_constant() def remove_useless_div(self): """Remove Div by 1""" div_nodes = [node for node in self.nodes() if node.op_type == "Div"] nodes_to_remove = [] for div in div_nodes: if self.find_constant_input(div, 1.0) == 1: nodes_to_remove.append(div) for node in nodes_to_remove: self.replace_input_of_all_nodes(node.output[0], node.input[0]) if nodes_to_remove: self.remove_nodes(nodes_to_remove) logger.info("Removed %d Div nodes", len(nodes_to_remove)) def convert_conv_to_nhwc(self): # Transpose weights in offline might help since ORT does not apply constant-folding on Transpose nodes. conv_to_nhwc_conv = FusionNhwcConv(self, update_weight=True) conv_to_nhwc_conv.apply() def merge_adjacent_transpose(self): fusion_transpose = FusionTranspose(self) fusion_transpose.apply() remove_count = 0 nodes = self.get_nodes_by_op_type("Transpose") for node in nodes: permutation = OnnxModel.get_node_attribute(node, "perm") assert isinstance(permutation, list) if permutation != list(range(len(permutation))): continue assert not ( self.find_graph_output(node.output[0]) or self.find_graph_input(node.input[0]) or self.find_graph_output(node.input[0]) ) # Let all children nodes skip current Transpose node and link to its parent # Note that we cannot update parent node output since parent node might have more than one children. self.replace_input_of_all_nodes(node.output[0], node.input[0]) self.remove_node(node) remove_count += 1 total = len(fusion_transpose.nodes_to_remove) + remove_count if total: logger.info("Removed %d Transpose nodes", total) def fuse_multi_head_attention(self, options: Optional[FusionOptions] = None): # Self Attention enable_packed_qkv = (options is None) or options.enable_packed_qkv self_attention_fusion = FusionAttentionUnet( self, self.hidden_size, self.num_heads, is_cross_attention=False, enable_packed_qkv=enable_packed_qkv, enable_packed_kv=False, ) self_attention_fusion.apply() # Cross Attention enable_packed_kv = (options is None) or options.enable_packed_kv cross_attention_fusion = FusionAttentionUnet( self, self.hidden_size, self.num_heads, is_cross_attention=True, enable_packed_qkv=False, enable_packed_kv=enable_packed_kv, ) cross_attention_fusion.apply() def fuse_bias_add(self): fusion = FusionBiasAdd(self) fusion.apply() def optimize(self, options: Optional[FusionOptions] = None): if is_installed("tqdm"): import tqdm from tqdm.contrib.logging import logging_redirect_tqdm with logging_redirect_tqdm(): steps = 18 progress_bar = tqdm.tqdm(range(steps), initial=0, desc="fusion") self._optimize(options, progress_bar) else: logger.info("tqdm is not installed. Run optimization without progress bar") self._optimize(options, None) def _optimize(self, options: Optional[FusionOptions] = None, progress_bar=None): if (options is not None) and not options.enable_shape_inference: self.disable_shape_inference() self.utils.remove_identity_nodes() if progress_bar: progress_bar.update(1) # Remove cast nodes that having same data type of input and output based on symbolic shape inference. self.utils.remove_useless_cast_nodes() if progress_bar: progress_bar.update(1) if (options is None) or options.enable_layer_norm: self.fuse_layer_norm() if progress_bar: progress_bar.update(1) if (options is None) or options.enable_gelu: self.fuse_gelu() if progress_bar: progress_bar.update(1) self.preprocess() if progress_bar: progress_bar.update(1) self.fuse_reshape() if progress_bar: progress_bar.update(1) if (options is None) or options.enable_group_norm: channels_last = (options is None) or options.group_norm_channels_last group_norm_fusion = FusionGroupNorm(self, channels_last) group_norm_fusion.apply() insert_transpose_fusion = FusionInsertTranspose(self) insert_transpose_fusion.apply() if progress_bar: progress_bar.update(1) if (options is None) or options.enable_bias_splitgelu: bias_split_gelu_fusion = FusionBiasSplitGelu(self) bias_split_gelu_fusion.apply() if progress_bar: progress_bar.update(1) if (options is None) or options.enable_attention: # self.save_model_to_file("before_mha.onnx") self.fuse_multi_head_attention(options) if progress_bar: progress_bar.update(1) if (options is None) or options.enable_skip_layer_norm: self.fuse_skip_layer_norm() if progress_bar: progress_bar.update(1) self.fuse_shape() if progress_bar: progress_bar.update(1) # Remove reshape nodes that having same shape of input and output based on symbolic shape inference. self.utils.remove_useless_reshape_nodes() if progress_bar: progress_bar.update(1) if (options is None) or options.enable_skip_group_norm: skip_group_norm_fusion = FusionSkipGroupNorm(self) skip_group_norm_fusion.apply() if progress_bar: progress_bar.update(1) if (options is None) or options.enable_bias_skip_layer_norm: # Fuse SkipLayerNormalization and Add Bias before it. self.fuse_add_bias_skip_layer_norm() if progress_bar: progress_bar.update(1) if options is not None and options.enable_gelu_approximation: self.gelu_approximation() if progress_bar: progress_bar.update(1) if options is None or options.enable_nhwc_conv: self.convert_conv_to_nhwc() self.merge_adjacent_transpose() if progress_bar: progress_bar.update(1) if options is not None and options.enable_bias_add: self.fuse_bias_add() if progress_bar: progress_bar.update(1) self.postprocess() if progress_bar: progress_bar.update(1) logger.info(f"opset version: {self.get_opset_version()}") def get_fused_operator_statistics(self): """ Returns node count of fused operators. """ op_count = {} ops = [ "Attention", "MultiHeadAttention", "LayerNormalization", "SkipLayerNormalization", "BiasSplitGelu", "GroupNorm", "SkipGroupNorm", "NhwcConv", "BiasAdd", ] for op in ops: nodes = self.get_nodes_by_op_type(op) op_count[op] = len(nodes) logger.info(f"Optimized operators:{op_count}") return op_count