# MIT License

#

# Copyright (c) Microsoft Corporation, Hugging Face. All rights reserved.

#

# Permission is hereby granted, free of charge, to any person obtaining a copy

# of this software and associated documentation files (the "Software"), to deal

# in the Software without restriction, including without limitation the rights

# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

# copies of the Software, and to permit persons to whom the Software is

# furnished to do so, subject to the following conditions:

#

# The above copyright notice and this permission notice shall be included in all

# copies or substantial portions of the Software.

#

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

# SOFTWARE.

from typing import Optional

import itertools

import numpy as np

import onnx

import packaging.version as pv

import warnings

from onnx import helper, numpy_helper

from onnx import onnx_pb as onnx_proto

import onnxslim.third_party.onnx_graphsurgeon as gs

FLOAT32 = 1

FLOAT16 = 10

def _npfloat16_to_int(np_list):

"""

Convert numpy float16 to python int.

:param np_list: numpy float16 list

:return int_list: python int list

"""

return [int(bin(_.view("H"))[2:].zfill(16), 2) for _ in np_list]

def convert_np_to_float16(np_array, min_positive_val=1e-7, max_finite_val=1e4):

"""

Convert float32 numpy array to float16 without changing sign or finiteness.

Positive values less than min_positive_val are mapped to min_positive_val.

Positive finite values greater than max_finite_val are mapped to max_finite_val.

Similar for negative values. NaN, 0, inf, and -inf are unchanged.

"""

def between(a, b, c):

return np.logical_and(a < b, b < c)

positive_values = np_array[np.where(np_array > 0)]

if positive_values.shape[0] > 0:

pos_max = positive_values.max()

pos_min = positive_values.min()

if pos_max >= max_finite_val:

warnings.warn(

"the float32 number {} will be truncated to {}".format(

pos_max, max_finite_val

)

)

if pos_min <= min_positive_val:

warnings.warn(

"the float32 number {} will be truncated to {}".format(

pos_min, min_positive_val

)

)

negative_values = np_array[np.where(np_array < 0)]

if negative_values.shape[0] > 0:

neg_max = negative_values.max()

neg_min = negative_values.min()

if neg_min <= -max_finite_val:

warnings.warn(

"the float32 number {} will be truncated to {}".format(

neg_min, -max_finite_val

)

)

if neg_max >= -min_positive_val:

warnings.warn(

"the float32 number {} will be truncated to {}".format(

neg_max, -min_positive_val

)

)

np_array = np.where(

between(0, np_array, min_positive_val), min_positive_val, np_array

)

np_array = np.where(

between(-min_positive_val, np_array, 0), -min_positive_val, np_array

)

np_array = np.where(

between(max_finite_val, np_array, float("inf")), max_finite_val, np_array

)

np_array = np.where(

between(float("-inf"), np_array, -max_finite_val), -max_finite_val, np_array

)

return np.float16(np_array)

def convert_tensor_float_to_float16(tensor, min_positive_val=1e-7, max_finite_val=1e4):

"""

Convert tensor float to float16.

:param tensor: TensorProto object

:return tensor_float16: converted TensorProto object

"""

if not isinstance(tensor, onnx_proto.TensorProto):

raise ValueError(

"Expected input type is an ONNX TensorProto but got %s" % type(tensor)

)

if tensor.data_type == onnx_proto.TensorProto.FLOAT:

tensor.data_type = onnx_proto.TensorProto.FLOAT16

# convert float_data (float type) to float16 and write to int32_data

if tensor.float_data:

float16_data = convert_np_to_float16(

np.array(tensor.float_data), min_positive_val, max_finite_val

)

int_list = _npfloat16_to_int(float16_data)

tensor.int32_data[:] = int_list

tensor.float_data[:] = []

# convert raw_data (bytes type)

if tensor.raw_data:

# convert n.raw_data to float

float32_list = np.fromstring(tensor.raw_data, dtype="float32")

# convert float to float16

float16_list = convert_np_to_float16(

float32_list, min_positive_val, max_finite_val

)

# convert float16 to bytes and write back to raw_data

tensor.raw_data = float16_list.tostring()

return tensor

def make_value_info_from_tensor(tensor):

shape = numpy_helper.to_array(tensor).shape

return helper.make_tensor_value_info(tensor.name, tensor.data_type, shape)

DEFAULT_OP_BLOCK_LIST = [

"ArrayFeatureExtractor",

"Binarizer",

"CastMap",

"CategoryMapper",

"DictVectorizer",

"FeatureVectorizer",

"Imputer",

"LabelEncoder",

"LinearClassifier",

"LinearRegressor",

"Normalizer",

"OneHotEncoder",

"RandomUniformLike",

"SVMClassifier",

"SVMRegressor",

"Scaler",

"TreeEnsembleClassifier",

"TreeEnsembleRegressor",

"ZipMap",

"NonMaxSuppression",

"TopK",

"RoiAlign",

"Resize",

# 'Range',

"CumSum",

"Min",

"Max",

"Upsample",

# NEW:

"RandomNormalLike",

# TODO: Ideally, "Cast" nodes should not be here, for the following reasons:

# - It breaks the semantics that the default list contains "ops that are not supported for float16 in ONNX Runtime".

# - When fp32 casts already exist in the model (e.g., for rotary embeddings), this script will insert redundant casts around it.

# However, without it, the graphs produced are invalid. Eventually, we will resolve this.

"Cast",

]

def initial_checking(model, disable_shape_infer):

func_infer_shape = None

if not disable_shape_infer and pv.Version(onnx.__version__) >= pv.Version("1.2"):

try:

from onnx.shape_inference import infer_shapes

func_infer_shape = infer_shapes

finally:

pass

if not isinstance(model, onnx_proto.ModelProto):

raise ValueError(

"Expected model type is an ONNX ModelProto but got %s" % type(model)

)

if func_infer_shape is not None:

model = func_infer_shape(model)

is_fp16_ready_flag = check_if_fp16_ready(model.graph)

return model, func_infer_shape, is_fp16_ready_flag

def convert_float_to_float16(

model,

min_positive_val=1e-7,

max_finite_val=1e4,

keep_io_types=False,

disable_shape_infer=False,

op_block_list=None,

node_block_list=None,

check_fp16_ready=True,

):

# create blocklists

if op_block_list is None:

op_block_list = DEFAULT_OP_BLOCK_LIST

if node_block_list is None:

node_block_list = []

op_block_list = set(op_block_list)

node_block_list = set(node_block_list)

global_input_name_dict = (

{}

) # key: input name, value: new output name after Cast node

# basic checking, including shape inference

model, func_infer_shape, is_fp16_ready_flag = initial_checking(

model, disable_shape_infer

)

if is_fp16_ready_flag and check_fp16_ready:

raise ValueError(

"The model is already converted to float16, if convert again, the model might be wrong. \n If you are sure to convert again, please set check_fp16_ready=False."

)

graph_stack = [model.graph]

is_top_level = True

while graph_stack:

next_level = []

for curr_graph in graph_stack:

process_graph_input(

curr_graph, is_top_level, keep_io_types, global_input_name_dict

)

value_info_block_list = process_tensor_in_node(

curr_graph,

op_block_list,

node_block_list,

min_positive_val,

max_finite_val,

)

process_value_info(curr_graph, value_info_block_list)

process_node_in_block_list(

curr_graph, global_input_name_dict, op_block_list, node_block_list

)

process_initializers(

curr_graph,

op_block_list,

node_block_list,

min_positive_val,

max_finite_val,

)

process_graph_output(curr_graph, is_top_level, keep_io_types)

sub_graph_list = get_next_level_graph(

curr_graph, op_block_list, node_block_list

)

if len(sub_graph_list) > 0:

next_level.extend(sub_graph_list)

if not is_top_level:

process_node_input_output(curr_graph, global_input_name_dict)

is_top_level = False # Going to process sub-graph

graph_stack = next_level

remove_unnecessary_cast_node(model.graph)

# Topologically sort the graph

# NOTE: We do not perform another round of optimization as the model is already optimized

graph = gs.import_onnx(model)

graph.toposort()

model = gs.export_onnx(graph)

return model

# Change the input/output of the node to the new output name after Cast node for sub-graph

# Because there have NO value_info start from

def process_node_input_output(

graph: onnx_proto.GraphProto, global_input_name_dict: dict

):

for node in graph.node:

for i, input_name in enumerate(node.input):

if input_name in global_input_name_dict:

node.input[i] = global_input_name_dict[input_name]

for i, output_name in enumerate(node.output):

if output_name in global_input_name_dict:

node.output[i] = global_input_name_dict[output_name]

def process_graph_input(

graph: onnx_proto.GraphProto,

is_top_level: bool,

is_io_fp32: bool,

global_input_name_dict: dict,

):

# The input dtype is float32, need to cast to fp16

if is_top_level and is_io_fp32:

for graph_input in graph.input: # n_input is ValueInfoProto

if graph_input.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:

downstream_nodes = find_downstream_node_by_input_name(

graph, graph_input.name

)

for d_node in downstream_nodes:

# More than one node may consume the model input, so we only create

# a single cast node, and then reuse this node when needed.

cast_exists = graph_input.name in global_input_name_dict

if cast_exists:

cast_node_output_name = global_input_name_dict[graph_input.name]

else:

cast_node_output_name = graph_input.name + "_fp16"

add_cast_node(

graph,

[graph_input.name],

[cast_node_output_name],

cast_node_output_name, # Set node name same as output name

FLOAT16,

)

add_new_value_info(

graph,

graph_input,

cast_node_output_name,

onnx_proto.TensorProto.FLOAT16,

)

for i, input_name in enumerate(d_node.input):

if input_name == graph_input.name:

d_node.input[i] = (

cast_node_output_name # Change the input of the second node

)

global_input_name_dict[graph_input.name] = (

cast_node_output_name

)

# For the sub-graph, don't do cast

else: # Change the input dtype to fp16 without any cast

for graph_input in graph.input:

if graph_input.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:

graph_input.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16

def process_graph_output(

graph: onnx_proto.GraphProto, is_top_level: bool, is_io_fp32: bool

):

if is_top_level and is_io_fp32: # the output dtype is float32, need to cast to fp16

for i, graph_output in enumerate(graph.output):

if graph_output.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:

new_producer_name = graph_output.name + "_fp16"

original_name = graph_output.name # The correct output name

# Get the node(s) that produce the model output

# These will most likely be fp16, but could be fp32 if the previous node is in block_list

upstream_nodes = find_upstream_node_by_output_name(graph, original_name)

assert len(upstream_nodes) == 1 # Should be only one node

producer_node = upstream_nodes[0]

for i, output_name in enumerate(producer_node.output):

if output_name == original_name:

producer_node.output[i] = new_producer_name

cast_node_name = new_producer_name + "_input_cast" + str(i)

add_cast_node(

graph,

[new_producer_name],

[original_name],

cast_node_name,

onnx_proto.TensorProto.FLOAT,

)

for value_info in graph.value_info:

if original_name == value_info.name:

value_info.type.tensor_type.elem_type = (

onnx_proto.TensorProto.FLOAT

)

# Get the node(s) that consume the model output

downstream_nodes = find_downstream_node_by_input_name(

graph,

original_name,

include_subgraphs=False,

)

# It is possible that the producer node is also input to downstream nodes

# So, we update the inputs of these downstream nodes

for d_node in downstream_nodes:

for i, input_name in enumerate(d_node.input):

if input_name == original_name:

d_node.input[i] = new_producer_name

else: # change the output dtype to fp16 in tensor

for graph_output in graph.output:

if graph_output.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:

graph_output.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16

def process_node_in_block_list(

graph: onnx_proto.GraphProto,

global_input_name_dict: dict,

op_block_list: list,

node_block_list: list,

):

# NB: Important to create a copy of the nodes in the graph to avoid modifying

# the graph in-place while iterating (causing an infinite loop)

for node in list(graph.node):

if (node.op_type in op_block_list) or (node.name in node_block_list):

insert_cast32_before_node(graph, node, global_input_name_dict)

insert_cast16_after_node(graph, node, global_input_name_dict)

# Todo: global_input_name_dict still not fill value

def insert_cast32_before_node(

graph: onnx_proto.GraphProto, node: onnx_proto.NodeProto, global_input_name_dict

):

for i, input_name in enumerate(node.input):

for value_info in itertools.chain(graph.value_info, graph.input):

if input_name == value_info.name:

if (

value_info.type.tensor_type.elem_type

!= onnx_proto.TensorProto.FLOAT16

):

break

cast_output_name = node.name + "_input_cast_" + str(i)

add_new_value_info(

graph, value_info, cast_output_name, onnx_proto.TensorProto.FLOAT

)

cast_node_name = node.name + "_input_cast" + str(i)

add_cast_node(

graph,

[input_name],

[cast_output_name],

cast_node_name,

onnx_proto.TensorProto.FLOAT,

)

node.input[i] = cast_output_name

break

# Todo: global_input_name_dict still not fill value

def insert_cast16_after_node(

graph: onnx_proto.GraphProto, node: onnx_proto.NodeProto, global_input_name_dict

):

for i, output_name in enumerate(node.output):

for value_info in itertools.chain(graph.value_info, graph.output):

if output_name == value_info.name:

if (

value_info.type.tensor_type.elem_type

!= onnx_proto.TensorProto.FLOAT

):

break

cast_input_name = node.name + "_output_cast_" + str(i)

add_new_value_info(

graph, value_info, cast_input_name, onnx_proto.TensorProto.FLOAT

)

value_info.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16

cast_node_name = node.name + "_output_cast" + str(i)

add_cast_node(

graph,

[cast_input_name],

[output_name],

cast_node_name,

onnx_proto.TensorProto.FLOAT16,

)

node.output[i] = cast_input_name

break

# Process tensor data in attribute of the node

def process_tensor_in_node(

graph: onnx_proto.GraphProto,

op_block_list: list,

node_block_list: list,

min_positive_val,

max_finite_val,

):

value_info_block_list = set() # This is for later use, not in this step

for node in graph.node:

# NOTE: "Cast" operation cannot change its output type because it is strongly typed.

if (

(node.op_type in op_block_list)

or (node.name in node_block_list)

or (node.op_type == "Cast")

):

# if (node.op_type in op_block_list) or (node.name in node_block_list):

# Only need to block the output value_info changing

for output_name in node.output:

value_info_block_list.add(output_name)

else:

for attr in node.attribute:

# one tensor

if attr.t.data_type == onnx_proto.TensorProto.FLOAT:

attr.t.CopyFrom(

convert_tensor_float_to_float16(

attr.t, min_positive_val, max_finite_val

)

)

# list of tensor

for t in attr.tensors:

if t.data_type == onnx_proto.TensorProto.FLOAT:

t.CopyFrom(

convert_tensor_float_to_float16(

t, min_positive_val, max_finite_val

)

)

return value_info_block_list

# Change all the value info type from float32 to float16 if not in block list

def process_value_info(graph: onnx_proto.GraphProto, value_info_block_list: list):

for value_info in graph.value_info:

if value_info.name in value_info_block_list:

continue

else:

if value_info.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:

value_info.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16

# Initializer is 'edge' type, so doesn't have value_info

def process_initializers(

graph: onnx_proto.GraphProto,

op_block_list,

node_block_list,

min_positive_val,

max_finite_val,

):

# Find the input of the block node, don't need to change this kind of initializer

initializer_block_list = set()

for node in graph.node:

if (node.op_type in op_block_list) or (node.name in node_block_list):

for (

input_name

) in (

node.input

): # some is initializer, some is value_info, can't distinguish but doesn't matter

initializer_block_list.add(input_name)

# Process initializers

for initializer in graph.initializer:

if initializer.name not in initializer_block_list:

if initializer.data_type == onnx_proto.TensorProto.FLOAT:

convert_tensor_float_to_float16(

initializer, min_positive_val, max_finite_val

)

def get_next_level_graph(

graph: onnx_proto.GraphProto, op_block_list: list, node_block_list: list

):

sub_graph_list = []

for node in graph.node:

if node.op_type in op_block_list or node.name in node_block_list:

continue

for attr in node.attribute:

# Check if sub-graph exist

if len(attr.g.node) > 0: # single sub-graph

sub_graph_list.append(attr.g)

for g in attr.graphs:

if len(g.node) > 0: # multiple sub-graphs

sub_graph_list.append(g)

return sub_graph_list

def add_cast_node(

graph: onnx_proto.GraphProto,

inputs: list,

outputs: list,

node_name: str,

to_type: int,

):

new_node = [helper.make_node("Cast", inputs, outputs, to=to_type, name=node_name)]

graph.node.extend(new_node)

def add_new_value_info(

graph: onnx_proto.GraphProto,

exist_value_info: onnx_proto.ValueInfoProto,

name: str,

dtype: int,

):

new_value_info = graph.value_info.add()

new_value_info.CopyFrom(exist_value_info)

new_value_info.name = name

new_value_info.type.tensor_type.elem_type = dtype

# Find the node that has the specified output name

def find_upstream_node_by_output_name(graph: onnx_proto.GraphProto, output_name: str):

nodes = []

for node in graph.node:

if output_name in node.output:

nodes.append(node)

assert len(nodes) <= 1 # Suppose there is less than one node found

return nodes

# Find the node that has the specified input name, including in subgraphs

def find_downstream_node_by_input_name(

graph: onnx_proto.GraphProto, input_name: str, include_subgraphs=True

):

nodes = []

# Check nodes in current graph

for node in graph.node:

if input_name in node.input:

nodes.append(node)

if not include_subgraphs:

continue

# Recursively check subgraphs in node attributes

for attr in node.attribute:

if attr.type == onnx_proto.AttributeProto.GRAPH:

# Single subgraph

if len(attr.g.node) > 0:

nodes.extend(find_downstream_node_by_input_name(attr.g, input_name))

# Multiple subgraphs

if attr.type == onnx_proto.AttributeProto.GRAPHS:

for g in attr.graphs:

if len(g.node) > 0:

nodes.extend(find_downstream_node_by_input_name(g, input_name))

return nodes

# Remove identity node

def remove_identity_node_from_model(model: onnx_proto.ModelProto):

remove_identity_node_from_graph(model.graph)

try:

from onnx.shape_inference import infer_shapes

func_infer_shape = infer_shapes

model = func_infer_shape(model)

return model

finally:

pass

# Remove identity node

def remove_identity_node_from_graph(graph: onnx_proto.GraphProto):

for curr_node in graph.node:

if curr_node.op_type == "Identity":

for input_name in curr_node.input:

upstream_nodes = find_upstream_node_by_output_name(graph, input_name)

for u_node in upstream_nodes:

if u_node is not None:

u_node.output[0] = curr_node.output[0]

graph.node.remove(curr_node)

def convert_float_to_float16_model_path(

model_path, min_positive_val=1e-7, max_finite_val=1e4, keep_io_types=False

):

"""

Convert tensor float type in the ONNX Model to tensor float16.

*It is to fix an issue that infer_shapes func cannot be used to infer >2GB models.

*But this function can be applied to all model sizes.

:param model_path: ONNX Model path

:return: converted ONNX ModelProto object

Examples

::

#Convert to ONNX ModelProto object and save model binary file:

from onnxmltools.utils.float16_converter import convert_float_to_float16_model_path

new_onnx_model = convert_float_to_float16_model_path('model.onnx')

onnx.save(new_onnx_model, 'new_model.onnx')

"""

disable_shape_infer = False

if pv.Version(onnx.__version__) >= pv.Version("1.8"):

try:

# infer_shapes_path can be applied to all model sizes

from onnx.shape_inference import infer_shapes_path

import tempfile

import os

# shape_infer_model_path should be in the same folder of model_path

with tempfile.NamedTemporaryFile(

dir=os.path.dirname(model_path)

) as tmpfile:

shape_infer_model_path = tmpfile.name

infer_shapes_path(model_path, shape_infer_model_path)

model = onnx.load(shape_infer_model_path)

disable_shape_infer = True

finally:

pass

if not disable_shape_infer:

model = onnx.load(model_path)

return convert_float_to_float16(

model, min_positive_val, max_finite_val, keep_io_types, disable_shape_infer

)

def remove_unnecessary_cast_node(graph_proto: onnx_proto.GraphProto):

# 1. find all cast nodes in the graph

cast_node_list = []

input_name_to_cast_node_dict = {}

output_name_to_cast_node_dict = {}

# using name as key to point to a node. because node object cannot be key

name_to_node_dict = {}

for node in graph_proto.node:

if node.op_type == "Cast":

# if node.name not in ["graph_input_cast0", "graph_output_cast0"]:

cast_node_list.append(node)

name_to_node_dict[node.name] = node

for input_name in node.input:

input_name_to_cast_node_dict[input_name] = node

for output_name in node.output:

output_name_to_cast_node_dict[output_name] = node

# 2. find upstream and downstream node of the cast node

cast_node_upstream_dict = {} # mapping cast node(name) to its upstream node

cast_node_downstream_dict = {} # mapping cast node(name) to its downstream node

for current_node in graph_proto.node:

# find the downstream node(s)

for input_name in current_node.input:

if input_name in output_name_to_cast_node_dict:

# found the downstream node of the cast node, might be multiple

cast_node = output_name_to_cast_node_dict[input_name]

if cast_node.name not in cast_node_downstream_dict:

cast_node_downstream_dict[cast_node.name] = current_node

else: # already exists one downstream node, make it a list

existing_downstream_nodes = cast_node_downstream_dict[

cast_node.name

]

if isinstance(existing_downstream_nodes, list):

existing_downstream_nodes.append(current_node)

else: # make a list

existing_downstream_nodes = [

existing_downstream_nodes,

current_node,

]

cast_node_downstream_dict[cast_node.name] = (

existing_downstream_nodes

)

# find the upstream node

for output_name in current_node.output:

if output_name in input_name_to_cast_node_dict:

# found the upstream node of the cast node, should be unique

cast_node = input_name_to_cast_node_dict[output_name]

cast_node_upstream_dict[cast_node.name] = current_node

# 3. remove the cast node which upstream is 'Constant'

for cast_node_name, upstream_node in cast_node_upstream_dict.items():

cast_node = name_to_node_dict[cast_node_name]

if upstream_node.op_type == "Constant":

cast_node_list.remove(cast_node)

# 4. find (cast_to_fp16, cast_to_fp32) pairs where --fp32--> cast_to_fp16 --fp16--> cast_to_fp32.

remove_candidate = []

name_to_value_info = {

value_info.name: value_info

for value_info in itertools.chain(graph_proto.value_info, graph_proto.input)

}

def get_type(name: str) -> Optional[int]:

if name in name_to_value_info:

return name_to_value_info[name].type

else:

# `name` has no value info.

return None

for cast_node_name, downstream_node in cast_node_downstream_dict.items():

cast_node = name_to_node_dict[cast_node_name]

if len(cast_node.input) != 1:

raise RuntimeError(

f"Cast node {cast_node_name} should have only one input, but has {len(cast_node.input)}."

)

input_type = get_type(cast_node.input[0])

if input_type != onnx_proto.TensorProto.FLOAT:

continue

if isinstance(downstream_node, list):

for dn in downstream_node:

if (

dn.op_type == "Cast"

and dn.attribute[0].i == 32

and cast_node.attribute[0].i == 16

and dn in cast_node_list

and cast_node in cast_node_list

):

remove_candidate.append((cast_node, dn))

else:

if (

downstream_node.op_type == "Cast"

and cast_node.attribute[0].i == FLOAT16

and downstream_node.attribute[0].i == FLOAT32

and downstream_node in cast_node_list

and cast_node in cast_node_list

):

remove_candidate.append((cast_node, downstream_node))

# 5. change "upstream --fp32--> cast_to_fp16 --fp16--> cast_to_fp32 --fp32--> downstream" to

# "upstream --fp32--> downstream".

for cast_node_pair in remove_candidate:

first_cast_node = cast_node_pair[0]

second_cast_node = cast_node_pair[1]

upstream_node = cast_node_upstream_dict.get(first_cast_node.name)

downstream_node = cast_node_downstream_dict.get(second_cast_node.name)

if upstream_node is None and downstream_node is not None:

# The upstream_node should be graph input

out = first_cast_node.input[0]

for i, input_name in enumerate(downstream_node.input):

for output_name in second_cast_node.output:

if input_name == output_name:

# change the input as the upstream node's output

downstream_node.input[i] = out

elif upstream_node is not None and downstream_node is None:

raise ValueError(

"The downstream node of the second cast node should be graph output"

)

else:

# find the upstream node's output to first_cast_node

out = None

for output_name in upstream_node.output:

if output_name == first_cast_node.input[0]:

out = output_name

break

# find the downstream node's input as second_cast_node's output

for i, input_name in enumerate(downstream_node.input):

for output_name in second_cast_node.output:

if input_name == output_name:

# change the input as the upstream node's output

downstream_node.input[i] = out

# 6. remove the cast node pair

for cast_node_pair in remove_candidate:

graph_proto.node.remove(cast_node_pair[0])

graph_proto.node.remove(cast_node_pair[1])

# Check if the model is already converted to float16

def check_if_fp16_ready(graph_proto):

# Check graph input and ouput

is_value_info_fp16 = False

for value_info in itertools.chain(

graph_proto.output, graph_proto.input, graph_proto.value_info

):

if value_info.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT16:

is_value_info_fp16 = True

break

# Check initializer

is_initializer_fp16 = False

for initializer in graph_proto.initializer:

if initializer.data_type == onnx_proto.TensorProto.FLOAT16:

is_initializer_fp16 = True

break

# Check cast node

has_cast_node_fp16 = False

for node in graph_proto.node:

if node.op_type == "Cast" and node.attribute[0].i == FLOAT16:

has_cast_node_fp16 = True

break

# Any of above flags is True, return True

if is_value_info_fp16 or is_initializer_fp16 or has_cast_node_fp16:

return True # already converted to float16

else:

return False # not converted to float16 yet