mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
017870a633
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23097 Test Plan: Imported from OSS Differential Revision: D16383329 Pulled By: suo fbshipit-source-id: 282f8bac2245d584b66139daf4e5ea7b2b317295
68 lines
2.3 KiB
C++
68 lines
2.3 KiB
C++
#pragma once
|
|
|
|
#include <torch/csrc/jit/ir.h>
|
|
#include <torch/csrc/jit/script/module.h>
|
|
|
|
#include <istream>
|
|
|
|
namespace caffe2 {
|
|
namespace serialize {
|
|
class ReadAdapterInterface;
|
|
} // namespace serialize
|
|
} // namespace caffe2
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
static script::ExtraFilesMap default_extra_files;
|
|
|
|
TORCH_API script::Module import_ir_module(
|
|
std::shared_ptr<script::CompilationUnit> cu,
|
|
const std::string& filename,
|
|
c10::optional<c10::Device> device = c10::nullopt,
|
|
script::ExtraFilesMap& extra_files = default_extra_files);
|
|
|
|
TORCH_API script::Module import_ir_module(
|
|
std::shared_ptr<script::CompilationUnit> cu,
|
|
std::istream& in,
|
|
c10::optional<c10::Device> device = c10::nullopt,
|
|
script::ExtraFilesMap& extra_files = default_extra_files);
|
|
|
|
TORCH_API script::Module import_ir_module(
|
|
std::shared_ptr<script::CompilationUnit> cu,
|
|
std::unique_ptr<caffe2::serialize::ReadAdapterInterface> rai,
|
|
c10::optional<c10::Device> device = c10::nullopt,
|
|
script::ExtraFilesMap& extra_files = default_extra_files);
|
|
|
|
/// Loads a serialized `script::Module` from the given `istream`.
|
|
///
|
|
/// The istream must contain a serialized `script::Module`, exported via
|
|
/// `torch::jit::ExportModule` in C++.
|
|
TORCH_API script::Module load(
|
|
std::istream& in,
|
|
c10::optional<c10::Device> device = c10::nullopt,
|
|
script::ExtraFilesMap& extra_files = default_extra_files);
|
|
|
|
/// Loads a serialized `script::Module` from the given `filename`.
|
|
///
|
|
/// The file stored at the location given in `filename` must contain a
|
|
/// serialized `script::Module`, exported either via `ScriptModule.save()` in
|
|
/// Python or `torch::jit::ExportModule` in C++.
|
|
TORCH_API script::Module load(
|
|
const std::string& filename,
|
|
c10::optional<c10::Device> device = c10::nullopt,
|
|
script::ExtraFilesMap& extra_files = default_extra_files);
|
|
|
|
/// Loads a serialized `script::Module` from the given `rai`.
|
|
///
|
|
/// The reader adapter, which is for customized input stream, must contain a
|
|
/// serialized `script::Module`, exported either via `ScriptModule.save()` in
|
|
/// Python or `torch::jit::ExportModule` in C++.
|
|
TORCH_API script::Module load(
|
|
std::unique_ptr<caffe2::serialize::ReadAdapterInterface> rai,
|
|
c10::optional<c10::Device> device = c10::nullopt,
|
|
script::ExtraFilesMap& extra_files = default_extra_files);
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|