ONNX-GraphSurgeon（简称GS）是一个用于创建、修改和分析ONNX模型的工具。ONNX（Open Neural Network Exchange）是一种开放的格式，它允许不同深度学习框架之间的模型互操作性。GS的核心是一个中间表示（IR），可以从ONNX导入并导出，通过简单的Python API，几乎可以对IR进行任何所需的修改。这包括更改节点、张量或图的属性，重新连接节点和张量，修改权重，添加、删除或替换子图。此外，GS还提供了额外的便利功能，如常量折叠、拓扑排序和删除死层。
在讲解如何使用ONNX-GraphSurgeon之前，我们先了解一下Polygraphy。Polygraphy是一个强大的工具集，它利用了ONNX-GraphSurgeon的功能，提供了Python API和命令行接口来验证、调试和优化ONNX模型。它使得用户能够轻松地探索模型的结构，执行各种分析任务，以及在需要时进行模型转换。
创建新模型在ONNX-GraphSurgeon中有两种方法：低级API和高级API。低级API直接暴露节点和张量，需要用户显式创建中间张量和节点，并将其添加到图中。例如，创建一个简单的加法操作：
```python
A = gs.Variable(name="A", shape=(1, -1), dtype=np.float32)
B = gs.Variable(name="B")
C = gs.Variable(name="C", shape=(1, -1), dtype=np.float32)
add0 = gs.Node(op="Add", inputs=[A, A], outputs=[B])
add1 = gs.Node(op="Add", inputs=[B, A], outputs=[C])
graph = gs.Graph(nodes=[add0, add1], inputs=[A], outputs=[C])
```
相比之下，高级API允许注册可重用的操作，自动处理节点和张量的创建。通过`Graph.register()`方法，用户可以定义自己的层，例如加法层：
```python
@gs.Graph.register()
def add(self, a, b):
return self.layer(op="Add", inputs=[a, b], outputs=["add_out_gs"])[0]
A = gs.Variable(name="A", shape=(1, -1), dtype=np.float32)
graph = gs.Graph(inputs=[A])
B = graph.add(A, A)
C = graph.add(B, A)
graph.outputs = [C]
```
在修改模型时，ONNX-GraphSurgeon提供了一种灵活的方法来遍历和修改图中的每个节点。例如，你可以遍历所有节点并应用某种修改，如更改操作类型或调整输入输出：
```python
for node in graph.nodes:
if node.op == "Add":
new_node = gs.Node(op="Multiply", inputs=node.inputs, outputs=node.outputs)
graph.replace_node(node, new_node)
```
TensorRT是NVIDIA开发的一个高性能的深度学习推理引擎，它可以将ONNX模型转换为优化的运行时图，以在GPU上实现更快的推理速度。Polygraphy和ONNX-GraphSurgeon可以帮助用户在转换到TensorRT之前准备和验证ONNX模型，确保其兼容性和性能。
总结起来，ONNX-GraphSurgeon是一个强大的工具，它提供了对ONNX模型的深入控制，使开发者能够轻松创建、修改和优化模型。结合Polygraphy，它为深度学习从业者提供了全面的模型分析和转换能力，对于跨框架工作流和优化推理性能尤其有用。无论是构建新模型还是优化现有模型，ONNX-GraphSurgeon都能提供必要的灵活性和便利性。
