소스 검색

修改剪枝后文件不能转换om问题

zhy 1 일 전
부모
커밋
00901aaac8
1개의 변경된 파일1개의 추가작업 그리고 1개의 파일을 삭제
  1. 1 1
      tests/prune_tool.py

+ 1 - 1
tests/prune_tool.py

@@ -43,7 +43,7 @@ def prune_weights(model_path, pruned_model, pruning_percentage=0.05):
     )
 
     # 创建新的模型
-    new_model = onnx.helper.make_model(new_graph, producer_name='onnx-example')
+    new_model = onnx.helper.make_model(new_graph, producer_name='onnx-example', opset_imports=model.opset_import)
     # 保存剪枝后的模型
     onnx.save(new_model, pruned_model)