1 | #pragma once |
2 | |
3 | #include <torch/csrc/jit/ir/ir.h> |
4 | #include <torch/csrc/onnx/onnx.h> |
5 | #include <unordered_map> |
6 | |
7 | namespace torch { |
8 | namespace jit { |
9 | |
10 | TORCH_API std::shared_ptr<Graph> ToONNX( |
11 | std::shared_ptr<Graph>& state, |
12 | ::torch::onnx::OperatorExportTypes operator_export_type); |
13 | TORCH_API std::unordered_map<Value*, Value*> BlockToONNX( |
14 | Block* old_block, |
15 | Block* new_block, |
16 | ::torch::onnx::OperatorExportTypes operator_export_type, |
17 | std::unordered_map<Value*, Value*>& env, |
18 | bool is_sub_block = false); |
19 | TORCH_API void NodeToONNX( |
20 | Node* old_node, |
21 | Block* new_block, |
22 | ::torch::onnx::OperatorExportTypes operator_export_type, |
23 | std::unordered_map<Value*, Value*>& env); |
24 | TORCH_API void RemovePrintOps(std::shared_ptr<Graph>& graph); |
25 | TORCH_API void PreprocessCaffe2Ops(std::shared_ptr<Graph>& graph); |
26 | |
27 | } // namespace jit |
28 | } // namespace torch |
29 | |