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
|
import torch
from torch.jit._serialization import validate_map_location
import pathlib
import os
def _load_for_lite_interpreter(f, map_location=None):
r"""
Load a :class:`LiteScriptModule`
saved with :func:`torch.jit._save_for_lite_interpreter`
Arguments:
f: a file-like object (has to implement read, readline, tell, and seek),
or a string containing a file name
map_location: a string or torch.device used to dynamically remap
storages to an alternative set of devices.
Returns:
A :class:`LiteScriptModule` object.
Example:
.. testcode::
import torch
import io
# Load LiteScriptModule from saved file path
torch.jit._load_for_lite_interpreter('lite_script_module.pt')
# Load LiteScriptModule from io.BytesIO object
with open('lite_script_module.pt', 'rb') as f:
buffer = io.BytesIO(f.read())
# Load all tensors to the original device
torch.jit.mobile._load_for_lite_interpreter(buffer)
"""
if isinstance(f, str):
if not os.path.exists(f):
raise ValueError("The provided filename {} does not exist".format(f))
if os.path.isdir(f):
raise ValueError("The provided filename {} is a directory".format(f))
map_location = validate_map_location(map_location)
if isinstance(f, str) or isinstance(f, pathlib.Path):
cpp_module = torch._C._load_for_lite_interpreter(f, map_location)
else:
cpp_module = torch._C._load_for_lite_interpreter_from_buffer(f.read(), map_location)
return LiteScriptModule(cpp_module)
class LiteScriptModule(object):
def __init__(self, cpp_module):
self._c = cpp_module
super(LiteScriptModule, self).__init__()
def __call__(self, *input):
return self._c.forward(input)
def find_method(self, method_name):
return self._c.find_method(method_name)
def forward(self, *input):
return self._c.forward(input)
def run_method(self, method_name, *input):
return self._c.run_method(method_name, input)
|
