Skip to content

Commit 68bda1b

Browse files
author
wanglichun
committed
update onnx.mdx
1 parent 0d3db37 commit 68bda1b

File tree

2 files changed

+21
-52
lines changed

2 files changed

+21
-52
lines changed

docs/faq/onnx.mdx

+21
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,26 @@ onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
7272
model_simp, check = onnx_simplifier.simplify(onnx_model, check_n = 0)
7373
onnx.save(model_simp, onnx_save_path)
7474
```
75+
76+
<details><summary>为了方便,针对这步我们提供了torchpipe.utils.onnx_export小工具</summary>
77+
78+
- 该工具可以实现将PyTorch模型转换为ONNX模型并保存到本地
79+
80+
```python
81+
import os
82+
from torchvision import models
83+
import torch
84+
import torchpipe
85+
86+
## export onnx
87+
m = models.resnet50(weights=None).eval()
88+
onnx_path = os.path.join("/tmp", f"resnet50.onnx")
89+
input = torch.randn(1, 3, 224, 224)
90+
torchpipe.utils.onnx_export(m, onnx_path, input, opset=17)
91+
92+
```
93+
94+
7595
### 转换失败说明
7696

7797

@@ -119,6 +139,7 @@ tgt_padding_mask = (tgt_in == self.eos_id).float()
119139

120140
## onnx 相关工具
121141

142+
122143
### [onnx-simplify](https://github.com/daquexian/onnx-simplifier) {#onnx-smi}
123144
简化模型结构的工具
124145
```python

docs/tools/onnx_export.mdx_

-52
This file was deleted.

0 commit comments

Comments
 (0)