def optimize_fp16_onnx_no_cast(input_onnx_path, optimized_onnx_path, epsilon): m = onnx.load(input_onnx_path) onnx_model = OnnxModel(m) nodes_to_remove = onnx_model.nodes() node_to_add = onnx.helper.make_node( "LayerNormalization", ["input", "layer_norm.weight", "layer_norm.bias"], ["output"], "layer_norm", epsilon=epsilon) onnx_model.remove_nodes(nodes_to_remove) onnx_model.add_node(node_to_add) onnx_model.prune_graph() onnx_model.save_model_to_file(optimized_onnx_path)
def optimize_fp16_onnx_no_cast(input_onnx_path, optimized_onnx_path, epsilon): m = onnx.load(input_onnx_path) onnx_model = OnnxModel(m) weight_name = get_weight(onnx_model) bias_name = get_bias(onnx_model) nodes_to_remove = [n for n in onnx_model.nodes() if n.output[0] != weight_name and n.output[0] != bias_name] nodes_to_remove = onnx_model.nodes() node_to_add = onnx.helper.make_node("LayerNormalization", ["input", weight_name, bias_name], ["output"], "layer_norm", epsilon=epsilon) onnx_model.remove_nodes(nodes_to_remove) onnx_model.add_node(node_to_add) onnx_model.prune_graph() onnx_model.save_model_to_file(optimized_onnx_path)