1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
#include <ATen/core/ivalue.h>
#include <caffe2/serialize/file_adapter.h>
#include <caffe2/serialize/inline_container.h>
#include <torch/csrc/jit/mobile/compatibility/backport.h>
#include <torch/csrc/jit/mobile/compatibility/backport_manager.h>
#include <torch/csrc/jit/mobile/compatibility/model_compatibility.h>
#include <string>
namespace torch {
namespace jit {
using caffe2::serialize::FileAdapter;
using caffe2::serialize::IStreamAdapter;
using caffe2::serialize::PyTorchStreamReader;
using caffe2::serialize::PyTorchStreamWriter;
using caffe2::serialize::ReadAdapterInterface;
const static BackportManager backportManager;
// Forward declare so that _backport_for_mobile() overloads can
// call this method directly.
bool _backport_for_mobile_impl(
std::istream& oss,
PyTorchStreamWriter& writer,
const int64_t to_version);
bool _backport_for_mobile(
std::istream& in,
std::ostream& out,
const int64_t to_version) {
auto writer_func = [&](const void* buf, size_t nbytes) -> size_t {
out.write(static_cast<const char*>(buf), nbytes);
return !out ? 0 : nbytes;
};
PyTorchStreamWriter writer(writer_func);
return _backport_for_mobile_impl(in, writer, to_version);
}
bool _backport_for_mobile(
std::istream& in,
const std::string& output_filename,
const int64_t to_version) {
PyTorchStreamWriter writer(output_filename);
return _backport_for_mobile_impl(in, writer, to_version);
}
bool _backport_for_mobile(
const std::string& input_filename,
std::ostream& out,
const int64_t to_version) {
std::ifstream file_stream;
std::unique_ptr<IStreamAdapter> istream_adapter;
file_stream.open(input_filename, std::ifstream::in | std::ifstream::binary);
if (!file_stream) {
AT_ERROR("open file failed, file path: ", input_filename);
}
auto writer_func = [&](const void* buf, size_t nbytes) -> size_t {
out.write(static_cast<const char*>(buf), nbytes);
return !out ? 0 : nbytes;
};
PyTorchStreamWriter writer(writer_func);
return _backport_for_mobile_impl(file_stream, writer, to_version);
}
bool _backport_for_mobile(
const std::string& input_filename,
const std::string& output_filename,
const int64_t to_version) {
std::ifstream file_stream;
file_stream.open(input_filename, std::ifstream::in | std::ifstream::binary);
if (!file_stream) {
AT_ERROR("open file failed, file path: ", input_filename);
}
PyTorchStreamWriter writer(output_filename);
return _backport_for_mobile_impl(file_stream, writer, to_version);
}
bool _backport_for_mobile_impl(
std::istream& oss,
PyTorchStreamWriter& writer,
const int64_t to_version) {
if (!backportManager.hasBytecodeBackportFunction(to_version + 1)) {
return false;
}
oss.seekg(0, oss.beg);
auto from_version = _get_model_bytecode_version(oss);
return backportManager.backport(oss, writer, from_version, to_version);
}
} // namespace jit
} // namespace torch
|
