算子-单算子ScatterElements

2 阅读1分钟
import onnx
from onnx import helper
from onnx import TensorProto

# 创建一个空的ONNX图
graph = helper.make_graph(nodes=[], name='Scatter_Graph', inputs=[], outputs=[])

# 创建Scatter节点的输入参数
data = helper.make_tensor_value_info('data', TensorProto.INT32, [3,3])  # 输入数据张量的形状为[1,5]
indices = helper.make_tensor_value_info('indices', TensorProto.INT32, [2,3])  # 索引张量的形状为[2, 2]
updates = helper.make_tensor_value_info('updates', TensorProto.INT32, [2,3])  # 更新数据张量的形状为[2, 2]
axis = 0  # 指定要进行散列操作的轴

graph.input.extend([data, indices, updates])

# 创建Scatter节点
scatter_node = helper.make_node('ScatterElements', ['data', 'indices', 'updates'], ['output'], name='Scatter_Node', axis=axis)

# 添加Scatter节点到图中
graph.node.extend([scatter_node])

# 创建Scatter节点的输出参数
output = helper.make_tensor_value_info('output', TensorProto.INT32, [3,3])  # 散列后的输出张量形状为[1,5]
graph.output.extend([output])

# 创建ONNX模型
model = helper.make_model(graph, producer_name='ONNX_Scatter_Elements_Demo')

# 保存ONNX模型到文件
onnx.save(model, 'scatter_elements_model3.onnx')