mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
1102f4c56e
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/19812 ghimport-source-id: 7cdb24c6a7501c6ec5f0eae07325512746f5abb9 Differential Revision: D15102803 Pulled By: zdevito fbshipit-source-id: bedf0bd6e1170fa294c65c87df75b82d8694f89c
55 lines
1.9 KiB
C++
55 lines
1.9 KiB
C++
#pragma once
|
|
|
|
#include <torch/csrc/jit/ir.h>
|
|
#include <torch/csrc/jit/script/module.h>
|
|
#include <torch/csrc/onnx/onnx.h>
|
|
|
|
#include <ostream>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
// This map is used to keep track of parameters that should be exported
|
|
// externally. When `defer_weight_export` is true, the returned map contains
|
|
// kv pairs that map {external reference name} -> {at::Tensor to be exported}.
|
|
// It is the responsibility of the caller to export these appropriately.
|
|
//
|
|
// For example, when exporting to a zip archive, the caller may write out files
|
|
// for each entry in the export map, with the filename being the key and the
|
|
// file contents being the raw tensor data.
|
|
using RawDataExportMap = std::unordered_map<std::string, at::Tensor>;
|
|
|
|
constexpr size_t CURRENT_OP_VERSION_SET = 1;
|
|
|
|
TORCH_API std::tuple<std::string, RawDataExportMap> export_onnx(
|
|
const std::shared_ptr<Graph>& graph,
|
|
const std::map<std::string, at::Tensor>& initializers,
|
|
int64_t onnx_opset_version,
|
|
bool defer_weight_export = false,
|
|
::torch::onnx::OperatorExportTypes operator_export_type =
|
|
::torch::onnx::OperatorExportTypes::ONNX,
|
|
bool strip_doc_string = true);
|
|
|
|
// For testing purposes
|
|
TORCH_API std::string pretty_print_onnx(
|
|
const std::shared_ptr<Graph>& graph,
|
|
const std::map<std::string, at::Tensor>& initializers,
|
|
int64_t onnx_opset_version,
|
|
bool defer_weight_export,
|
|
::torch::onnx::OperatorExportTypes operator_export_type =
|
|
::torch::onnx::OperatorExportTypes::ONNX,
|
|
bool google_printer = false);
|
|
|
|
TORCH_API void ExportModule(
|
|
const script::Module& module,
|
|
std::ostream& out,
|
|
const script::ExtraFilesMap& metadata = script::ExtraFilesMap());
|
|
|
|
TORCH_API void ExportModule(
|
|
const script::Module& module,
|
|
const std::string& filename,
|
|
const script::ExtraFilesMap& metadata = script::ExtraFilesMap());
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|