他的回复:
通过直接修改onnx模型,消除了警告,但是生成的om文件仍然带有 _ubuntu_x86_64 后缀。 修改onnx的文件如下: from opcode import opname from typing import List, Optional import onnx import numpy as np from onnx import helper, TensorProto, numpy_helper, checker from typing import List, Optional def insert_cast_for_op_inputs( model_path: str, output_path: str, target_op_names: list, # 指定的算子名称列表 cast_to_type: int = TensorProto.INT32 # 转换目标类型 ): """ 对指定算子列表中每个算子的所有int64输入插入Cast转为int32 """ # 1. 加载模型 model = onnx.load(model_path) graph = model.graph nodes = graph.node # 2. 遍历每个目标算子名称 for target_op_name in target_op_names: # 定位目标算子 target_node = None target_node_idx = -1 for idx, node in enumerate(nodes): if node.name == target_op_name: target_node = node target_node_idx = idx break if target_node is None: print(f"警告:未找到算子{target_op_name},跳过处理") continue else: print(f"开始处理算子{target_op_name}") # 3. 遍历目标算子的所有输入,逐个处理 cast_nodes_to_insert = [] # 存储需要插入的Cast节点 input_mapping = {} # 输入名→Cast输出名的映射 for input_idx, input_name in enumerate(target_node.input): # # 获取输入张量类型 if input_name == '/pt2pl_layers.0/aggr_module/Mod_output_0': # for idx, node in enumerate(nodes): # if node.output and node.output[0] == '/pt2pl_layers.0/aggr_module/Mod_output_0': # input_mapping[input_idx] = node # break input_mapping[input_idx] = '/pt2pl_layers.0/aggr_module/Mod_output_0' continue # 生成唯一的Cast输出名(算子名+输入名+后缀) cast_output_name = f"{target_op_name}_{input_name}_cast_to_{cast_to_type}" # 创建Cast节点 cast_node = helper.make_node( "Cast", inputs=[input_name], outputs=[cast_output_name], to=cast_to_type, name=f"Cast_{target_op_name}_{input_name}" ) cast_nodes_to_insert.append(cast_node) input_mapping[input_idx] = cast_output_name # 4. 替换目标算子的输入(将原输入替换为Cast输出) for input_idx, cast_output in input_mapping.items(): target_node.input[input_idx] = cast_output # 5. 插入所有Cast节点到目标算子之前(按输入顺序) # 反向插入,保证顺序与输入一致(因为list.insert是前插) for cast_node in reversed(cast_nodes_to_insert): nodes.insert(target_node_idx, cast_node) # 6. 验证模型合法性 try: checker.check_model(model) print("模型验证通过!") except checker.ValidationError as e: print(f"模型验证失败:{e}") return # 7. 保存修改后的模型 onnx.save(model, output_path) print(f"修改后的模型已保存至:{output_path}") def delete_op(onnx_path, op_list: List[str]): """ 删除 ONNX 模型中的指定算子 Args: onnx_path: ONNX 模型路径 op_list: 要删除的算子类型列表 """ # 加载模型 model = onnx.load(onnx_path) graph = model.graph # 保存原始节点数量 original_node_count = len(graph.node) # 过滤掉要删除的算子 filtered_nodes = [] deleted_nodes = [] for node in graph.node: if node.name in op_list: deleted_nodes.append(node) print(f"删除节点: {node.name}, 类型: {node.op_type}") # 重新连接被删除节点的输入输出 # 将下游节点中对本节点输出的引用替换为输入 input_tensor = node.input[0] if node.input else None output_tensor = node.output[0] if node.output else None if input_tensor and output_tensor: for other_node in graph.node: if other_node != node: # 不处理当前节点 # 替换输入 for i, inp in enumerate(other_node.input): if inp == output_tensor: other_node.input[i] = input_tensor # 替换输出(如果有节点以本节点输出作为其输出) for i, out in enumerate(other_node.output): if out == output_tensor: other_node.output[i] = input_tensor else: filtered_nodes.append(node) # 更新节点列表 graph.ClearField("node") graph.node.extend(filtered_nodes) # 清理未使用的值信息 used_tensors = set() # 收集所有使用的张量 for node in graph.node: used_tensors.update(node.input) used_tensors.update(node.output) # 清理未使用的值信息 filtered_value_info = [] for vi in graph.value_info: if vi.name in used_tensors: filtered_value_info.append(vi) graph.ClearField("value_info") graph.value_info.extend(filtered_value_info) # 清理初始值 filtered_initializer = [] for init in graph.initializer: if init.name in used_tensors: filtered_initializer.append(init) graph.ClearField("initializer") graph.initializer.extend(filtered_initializer) # 保存修改后的模型 output_path = onnx_path.replace(".onnx", "-cast.onnx") onnx.save(model, output_path) print(f"\n删除完成!") print(f"原始节点数: {original_node_count}") print(f"当前节点数: {len(graph.node)}") print(f"删除节点数: {len(deleted_nodes)}") print(f"修改后模型已保存到: {output_path}") return model, output_path def convert_specific_constants_int64_to_int32(input_model_path, output_model_path): """ 读取ONNX模型,查找指定的Constant节点,如果其数据类型为INT64, 则将其数据直接转换为INT32并修改模型属性。 """ # 1. 加载模型 model = onnx.load(input_model_path) graph = model.graph # 定义你需要修改的目标节点名称(基于你提供的日志) # 注意:脚本会根据名称匹配,但只会修改真正是INT64的节点 target_node_names = { # "/Constant_34", "/pt2pl_layers.0/Constant_34", "/pl2pl_layers.0/Constant_34", '/pt2pl_layers.0/Constant_20', '/pl2pl_layers.0/Constant_20', '/pt2pl_layers.0/Constant_25', '/pt2pl_layers.0/Constant_26', '/pt2pl_layers.0/Constant_27', '/pt2pl_layers.0/Constant_28', '/pl2pl_layers.0/Constant_25', '/pl2pl_layers.0/Constant_26', '/pl2pl_layers.0/Constant_27', '/pl2pl_layers.0/Constant_28', '/pt2pl_layers.0/Constant_29', '/pl2pl_layers.0/Constant_29', } modified_count = 0 print("开始扫描并修改指定节点...") for node in graph.node: # 仅处理 Constant 类型且在目标列表中的节点 if node.op_type == "Constant" and node.name in target_node_names: # 获取 Constant 节点的 value 属性 for attr in node.attribute: if attr.name == "value": tensor_proto = attr.t # 【关键步骤】检查数据类型 # ONNX Data Types: 1=FLOAT, 7=INT64, 6=INT32 if tensor_proto.data_type == TensorProto.INT64: print(f" -> 正在修改节点: {node.name} (INT64 -> INT32)") # 1. 使用 numpy_helper 提取原始数据 (自动处理 raw_data 或 int64_data) data_array = numpy_helper.to_array(tensor_proto) # 2. 强制转换为 int32 data_array_int32 = data_array.astype(np.int32) # 3. 创建新的 TensorProto (保留原有名称) new_tensor_proto = numpy_helper.from_array(data_array_int32, name=tensor_proto.name) # 4. 覆盖旧的 TensorProto 属性 # CopyFrom 会清空原有字段并复制新内容,确保 raw_data 和 int64_data 不会冲突 attr.t.CopyFrom(new_tensor_proto) modified_count += 1 else: # 提示用户:虽然节点在列表中,但因为不是INT64所以跳过(例如那个Float类型的Constant_34) print(f" -> 跳过节点: {node.name} (当前类型为 {tensor_proto.data_type}, 非INT64)") # 2. 保存修改后的模型 if modified_count > 0: onnx.save(model, output_model_path) print(f"\n修改完成。共修改 {modified_count} 个节点。") print(f"模型已保存至: {output_model_path}") else: print("\n未发现需要修改的INT64节点。") def convert_topk_indices_to_int32(input_model_path, output_model_path): """ 遍历ONNX模型,将所有TopK算子的索引输出类型从INT64 (7) 转换为 INT32 (6)。 """ try: model = onnx.load(input_model_path) graph = model.graph except Exception as e: print(f"加载模型失败: {e}") return modified_count = 0 # 1. 遍历并修改所有 TopK 节点的输出类型 for node in graph.node: if node.op_type == "TopK": print(node.name) # TopK 的第二个输出是索引 (Indices) index_output_name = node.output[1] # 2. 修改 ValueInfoProto 中的类型 found_value_info = False for value_info in graph.value_info: if value_info.name == index_output_name: if value_info.type.tensor_type.elem_type == TensorProto.INT64: value_info.type.tensor_type.elem_type = TensorProto.INT32 print(f" -> 修改 TopK 节点 '{node.name}' 的索引输出 '{index_output_name}' 类型为 INT32。") modified_count += 1 found_value_info = True break elif value_info.type.tensor_type.elem_type == TensorProto.INT32: # 已经是 INT32,无需修改 found_value_info = True break # 3. 检查并修改 Graph 输出列表中的类型 (如果 TopK 的索引是模型的最终输出) for output in graph.output: if output.name == index_output_name: if output.type.tensor_type.elem_type == TensorProto.INT64: output.type.tensor_type.elem_type = TensorProto.INT32 print(f" -> 修改模型输出 '{index_output_name}' 的类型为 INT32。") modified_count += 1 break # 4. 检查模型有效性并保存 if modified_count > 0: # 必须重新检查模型有效性,以确保修改是合法的 try: onnx.checker.check_model(model) onnx.save(model, output_model_path) print(f"\n修改完成。共修改 {modified_count} 处索引类型。") print(f"模型已保存至: {output_model_path}") except Exception as e: print(f"\n警告: 保存前模型检查失败。请检查修改是否导致图结构错误: {e}") else: print("\n未发现需要修改的 TopK INT64 索引。") def insert_cast_before_concat_input(input_model_path, output_model_path, concat_node_name, input_index_to_cast, cast_to=TensorProto.INT32): """ 在指定的 Concat 节点的一个输入前插入一个 Cast 算子。 Args: input_model_path (str): 原始 ONNX 模型路径。 output_model_path (str): 修改后的 ONNX 模型保存路径。 concat_node_name (str): 目标 Concat 算子的名称 (例如: "/pt2pl_layers.0/Concat_1")。 input_index_to_cast (int): 需要插入 Cast 算子的输入索引 (0-based, 例如第二个输入是 1)。 """ try: model = onnx.load(input_model_path) graph = model.graph except Exception as e: print(f"加载模型失败: {e}") return # 1. 查找目标 Concat 节点 concat_node = None for node in graph.node: if node.name == concat_node_name: concat_node = node break if not concat_node: print(f"错误: 在模型中未找到指定的 Concat 节点 '{concat_node_name}'。") return if input_index_to_cast >= len(concat_node.input): print(f"错误: 目标输入索引 {input_index_to_cast} 超出 Concat 节点输入数量 {len(concat_node.input)}。") return # 2. 确定 Cast 算子的输入、输出名称和目标类型 original_input_name = concat_node.input[input_index_to_cast] # 例如: x0:DT_FLOAT # 构造新的 Cast 算子的名称和输出名称 cast_node_name = f"Cast_to_INT32_for_{concat_node_name}_Input{input_index_to_cast}" new_input_name_to_concat = f"{cast_node_name}_output" # Cast 算子的输出,将作为 Concat 的新输入 # 3. 创建 Cast 算子 cast_node = helper.make_node( 'Cast', # op_type [original_input_name], # inputs (原始输入) [new_input_name_to_concat], # outputs (Cast 的新输出) name=cast_node_name, # node name to=cast_to # attribute: 目标数据类型 INT32 ) # 4. 修改 Concat 节点的输入 # 将 Concat 节点的第 input_index_to_cast 个输入从 original_input_name 替换为 new_input_name_to_concat print(f"修改前 Concat 输入列表: {list(concat_node.input)}") del concat_node.input[input_index_to_cast] concat_node.input.insert(input_index_to_cast, new_input_name_to_concat) print(f"修改后 Concat 输入列表: {list(concat_node.input)}") # 5. 将新的 Cast 节点插入到图的节点列表中 # 查找 Concat 节点在图中的位置,然后在其前面插入 Cast 节点 # 找到 Concat 节点在 graph.node 中的索引 concat_index = -1 for i, node in enumerate(graph.node): if node.name == concat_node_name: concat_index = i break # 在 Concat 节点前插入 Cast 节点 if concat_index != -1: graph.node.insert(concat_index, cast_node) print(f"\n成功插入 Cast 节点 '{cast_node_name}'。") print(f"Cast 节点插入在 '{concat_node_name}' 之前。") else: # 这应该不会发生,因为前面已经找到了 concat_node print("警告: 未能找到 Concat 节点在图中的位置。") # 6. 保存修改后的模型 try: onnx.checker.check_model(model) onnx.save(model, output_model_path) print(f"\n修改完成。新模型已保存至: {output_model_path}") except Exception as e: print(f"\n错误: 保存修改后的模型时发生问题,请检查图结构: {e}") # ------------------- 示例调用 ------------------- if __name__ == "__main__": remove_op_list = [ # '/pt2pl_layers.0/Cast_2', # '/pl2pl_layers.0/Cast_2', '/pt2pl_layers.0/Cast_5', '/pl2pl_layers.0/Cast_5', '/pt2pl_layers.0/Cast_6', '/pl2pl_layers.0/Cast_6', '/pt2pl_layers.0/Cast_7', '/pl2pl_layers.0/Cast_7', '/pt2pl_layers.0/Cast_4', '/pl2pl_layers.0/Cast_4', '/pt2pl_layers.0/Cast_8', '/pl2pl_layers.0/Cast_8', ] delete_op('batch_map_encoder.onnx', remove_op_list) convert_specific_constants_int64_to_int32('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx') convert_topk_indices_to_int32('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx') insert_cast_before_concat_input('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx', "/pt2pl_layers.0/Concat_1", 0) insert_cast_before_concat_input('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx', "/pl2pl_layers.0/Concat_1", 0) insert_cast_before_concat_input('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx', "/pt2pl_layers.0/TopK_1", 1) insert_cast_before_concat_input('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx', "/pl2pl_layers.0/TopK_1", 1) insert_cast_before_concat_input('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx', "/pt2pl_layers.0/ReduceSum_1", 0, TensorProto.FLOAT) insert_cast_before_concat_input('batch_map_encoder-cast.onnx','batch_map_encoder-cast.onnx', "/pl2pl_layers.0/ReduceSum_1", 0, TensorProto.FLOAT)