feat(add layernorm op): layernorm operator added

添加layernorm 算子

ONNX_Operator_Vis/LayerNorm/LayerNorm_export_onnx.py

@@ -0,0 +1,49 @@
+import torch.onnx
+import torch
+import onnx
+import onnxsim
+import torch.nn as nn
+
+
+class Model(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1)
+ self.norm = nn.LayerNorm(3)
+ self.act = nn.ReLU()
+
+ def forward(self, x):
+ _, _, H, W = x.shape
+ L = H * W
+ x = self.conv1(x)
+ x = x.view(x.shape[0], x.shape[1], L).permute(0, 2, 1)
+ x = self.norm(x)
+ x = self.act(x)
+ return x
+
+def export_onnx_graph():
+ input = torch.Tensor(1, 3, 5, 5).uniform_(-1, 1)
+ model = Model()
+ model.eval()
+
+ file = "./sample-ln-before.onnx"
+ torch.onnx.export(
+ model = model,
+ args = (input,),
+ f = file,
+ input_names = ["input0"],
+ output_names = ["output0"],
+ opset_version = 12)
+
+ print("\nFinished export {}".format(file))
+
+ model_onnx = onnx.load(file)
+ onnx.checker.check_model(model_onnx)
+
+ print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")
+ model_onnx, check = onnxsim.simplify(model_onnx)
+ assert check, "assert check failed"
+ onnx.save(model_onnx, file)
+
+
+export_onnx_graph()

ONNX_Operator_Vis/LayerNorm/layerNorm.py

@@ -28,7 +28,7 @@ def forward(self, x):
 # input_data = torch.randn(20, 100, 35, 45) # 输入形状为(1, 1, 5, 5)
 
 # 导出模型到ONNX格式
-onnx_path = "/root/bruce_cui/onnx_operator_vis/ONNX_Operators/layernorm.onnx"
-torch.onnx.export(model, input_data, onnx_path, verbose=True, opset_version=11)
+onnx_path = "/Users/bruce/Downloads/Chip_test_models/layernorm.onnx"
+torch.onnx.export(model, input_data, onnx_path, verbose=False, opset_version=11)
 
 print("ONNX模型已导出到:", onnx_path)

ONNX_Operator_Vis/LayerNorm/onnx_graphsurgon.py

@@ -0,0 +1,36 @@
+import onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+
+# Register functions to make graph generation easier
+@gs.Graph.register()
+def min(self, *args):
+ return self.layer(op="Min", inputs=args, outputs=["min_out"])[0]
+
+@gs.Graph.register()
+def max(self, *args):
+ return self.layer(op="Max", inputs=args, outputs=["max_out"])[0]
+
+@gs.Graph.register()
+def identity(self, inp):
+ return self.layer(op="Identity", inputs=[inp], outputs=["identity_out"])[0]
+
+
+# Generate the graph
+graph = gs.Graph()
+
+graph.inputs = [gs.Variable("input", shape=(4, 4), dtype=np.float32)]
+
+# Clip values to [0, 6]
+MIN_VAL = np.array(0, np.float32)
+MAX_VAL = np.array(6, np.float32)
+
+# Add identity nodes to make the graph structure a bit more interesting
+inp = graph.identity(graph.inputs[0])
+max_out = graph.max(graph.min(inp, MAX_VAL), MIN_VAL)
+graph.outputs = [graph.identity(max_out), ]
+
+# Graph outputs must include dtype information
+graph.outputs[0].to_variable(dtype=np.float32, shape=(4, 4))
+
+onnx.save(gs.export_onnx(graph), "model.onnx")

Tools/onnx_model_convert_onnxSimplify.py

@@ -18,4 +18,4 @@ def onnx_simplify(path):
 
 
 if __name__=="__main__":
- onnx_simplify("/Users/bruce/Downloads/Chip_test_models/models/yolov5_model/yolov5m.onnx")
+ onnx_simplify("/Users/bruce/Downloads/Chip_test_models/yolov5_modes_experiments/swim-transformer.onnx")