mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
5ec1c293eb
Differential Revision: D16552212 Original commit changeset: 7c7de5a096ad fbshipit-source-id: b5ea5f626883e2b213b9d02875e83e64ed206e58
450 lines
14 KiB
Python
450 lines
14 KiB
Python
"""
|
|
The weak_script annotation needs to be here instead of inside torch/jit/ so it
|
|
can be used in other places in torch/ (namely torch.nn) without running into
|
|
circular dependency problems
|
|
"""
|
|
|
|
import inspect
|
|
import weakref
|
|
import torch._C
|
|
from torch._six import builtins
|
|
|
|
# Wrapper functions that can call either of 2 functions depending on a boolean
|
|
# argument
|
|
boolean_dispatched = weakref.WeakKeyDictionary() # noqa: T484
|
|
|
|
|
|
def createResolutionCallback(frames_up=0):
|
|
"""
|
|
Creates a function which, given a string variable name,
|
|
returns the value of the variable in the scope of the caller of
|
|
the function which called createResolutionCallback (by default).
|
|
|
|
This is used to enable access in-scope Python variables inside
|
|
TorchScript fragments.
|
|
|
|
frames_up is number of additional frames to go up on the stack.
|
|
The default value is 0, which correspond to the frame of the caller
|
|
of createResolutionCallback. Also for example, if frames_up is set
|
|
to 1, then the frame of the caller's caller of createResolutionCallback
|
|
will be taken.
|
|
|
|
For example, the following program prints 2::
|
|
|
|
def bar():
|
|
cb = createResolutionCallback(1)
|
|
print(cb("foo"))
|
|
|
|
def baz():
|
|
foo = 2
|
|
bar()
|
|
|
|
baz()
|
|
"""
|
|
frame = inspect.currentframe()
|
|
i = 0
|
|
while i < frames_up + 1:
|
|
frame = frame.f_back
|
|
i += 1
|
|
|
|
f_locals = frame.f_locals
|
|
f_globals = frame.f_globals
|
|
|
|
def env(key):
|
|
if key in f_locals:
|
|
return f_locals[key]
|
|
elif key in f_globals:
|
|
return f_globals[key]
|
|
elif hasattr(builtins, key):
|
|
return getattr(builtins, key)
|
|
|
|
return env
|
|
|
|
|
|
def get_closure(fn):
|
|
"""
|
|
Get a dictionary of closed over variables from a function
|
|
"""
|
|
captures = {}
|
|
captures.update(fn.__globals__)
|
|
|
|
for index, captured_name in enumerate(fn.__code__.co_freevars):
|
|
captures[captured_name] = fn.__closure__[index].cell_contents
|
|
|
|
return captures
|
|
|
|
|
|
def createResolutionCallbackFromClosure(fn):
|
|
"""
|
|
Create a resolutionCallback by introspecting the function instead of
|
|
looking up the stack for the enclosing scope
|
|
"""
|
|
closure = get_closure(fn)
|
|
|
|
def env(key):
|
|
if key in closure:
|
|
return closure[key]
|
|
elif hasattr(builtins, key):
|
|
return getattr(builtins, key)
|
|
return None
|
|
|
|
return env
|
|
|
|
|
|
def can_compile_class(cls):
|
|
# If any of the functions on a type don't have a code object, this type can't
|
|
# be compiled and is probably a builtin / bound from C
|
|
fns = [getattr(cls, name) for name in cls.__dict__ if inspect.isroutine(getattr(cls, name))]
|
|
has_code = [hasattr(fn, '__code__') for fn in fns]
|
|
return all(has_code)
|
|
|
|
|
|
def createResolutionCallbackForClassMethods(cls):
|
|
"""
|
|
This looks at all the methods defined in a class and pulls their closed-over
|
|
variables into a dictionary and uses that to resolve variables.
|
|
"""
|
|
# cls is a type here, so `ismethod` is false since the methods on the type
|
|
# aren't bound to anything, so Python treats them as regular functions
|
|
fns = [getattr(cls, name) for name in cls.__dict__ if inspect.isroutine(getattr(cls, name))]
|
|
captures = {}
|
|
|
|
for fn in fns:
|
|
captures.update(get_closure(fn))
|
|
|
|
return lambda key: captures.get(key, None)
|
|
|
|
|
|
def boolean_dispatch(arg_name, arg_index, default, if_true, if_false, module_name, func_name):
|
|
"""
|
|
Dispatches to either of 2 script functions based on a boolean argument.
|
|
In TorchScript, the boolean argument must be constant so that the correct
|
|
function to use can be determined at compile time.
|
|
"""
|
|
def fn(*args, **kwargs):
|
|
dispatch_flag = False
|
|
if arg_name in kwargs:
|
|
dispatch_flag = kwargs[arg_name]
|
|
elif arg_index < len(args):
|
|
dispatch_flag = args[arg_index]
|
|
|
|
if dispatch_flag:
|
|
return if_true(*args, **kwargs)
|
|
else:
|
|
return if_false(*args, **kwargs)
|
|
|
|
if if_true.__doc__ is None and if_false.__doc__ is not None:
|
|
doc = if_false.__doc__
|
|
if_true.__doc__ = doc
|
|
elif if_false.__doc__ is None and if_true.__doc__ is not None:
|
|
doc = if_true.__doc__
|
|
if_false.__doc__ = doc
|
|
elif if_false.__doc__ is None and if_true.__doc__ is None:
|
|
# neither function has a docstring
|
|
doc = None
|
|
else:
|
|
raise RuntimeError("only one function can have a docstring")
|
|
fn.__doc__ = doc
|
|
|
|
if module_name is not None:
|
|
fn.__module__ = module_name
|
|
if func_name is not None:
|
|
fn.__name__ = func_name
|
|
|
|
boolean_dispatched[fn] = {
|
|
"if_true": if_true,
|
|
"if_false": if_false,
|
|
"index": arg_index,
|
|
"default": default,
|
|
"arg_name": arg_name
|
|
}
|
|
return fn
|
|
|
|
|
|
|
|
class FunctionModifiers(object):
|
|
"""
|
|
Used to denote the behavior of a function in TorchScript. See export() and
|
|
ignore() for details.
|
|
"""
|
|
IGNORE_AND_DROP = "ignore (leave as a call to Python, replace with a 'raise' on torch.jit.save)"
|
|
IGNORE = "ignore (leave as a call to Python, cannot be torch.jit.save'd)"
|
|
EXPORT = "export (compile this function even if nothing calls it)"
|
|
DEFAULT = "default (compile if called from a exported function / forward)"
|
|
|
|
|
|
def export(fn):
|
|
"""
|
|
This decorator indicates that a method is used as an entry point into a
|
|
ScriptModule. `forward` implicitly is used as an entry point, so it does
|
|
not need this decorator.
|
|
|
|
Methods are added to a ScriptModule as they are called in Python. If a
|
|
method is never called, it will not be included in the ScriptModule when
|
|
saving. This decorator explicitly marks that a method should be included
|
|
even if it is not called from Python.
|
|
"""
|
|
fn._torchscript_modifier = FunctionModifiers.EXPORT
|
|
return fn
|
|
|
|
|
|
def ignore(drop_on_export=False):
|
|
"""
|
|
This decorator indicates to the compiler that a function or method should
|
|
be ignored and left as a Python function.
|
|
|
|
With `drop_on_export=False` (the default), calls to this function will
|
|
prevent saving a TorchScript model.
|
|
|
|
With `drop_on_export=True`, any calls to this function from other
|
|
TorchScript code will be replaced with a `raise`. This allows you to leave
|
|
code in your TorchScript model that is only ever run when the Python
|
|
interpreter is present.
|
|
"""
|
|
if callable(drop_on_export):
|
|
# used without any args, so drop_on_export is actually a function
|
|
# @torch.jit.ignore
|
|
# def fn(...):
|
|
fn = drop_on_export
|
|
fn._torchscript_modifier = FunctionModifiers.IGNORE
|
|
return fn
|
|
|
|
if isinstance(drop_on_export, bool):
|
|
def decorator(fn):
|
|
if drop_on_export:
|
|
fn._torchscript_modifier = FunctionModifiers.IGNORE_AND_DROP
|
|
else:
|
|
fn._torchscript_modifier = FunctionModifiers.IGNORE
|
|
return fn
|
|
return decorator
|
|
raise RuntimeError("Argument to @torch.jit.ignore must be a bool or "
|
|
"a function but got {}".format(drop_on_export))
|
|
|
|
|
|
def should_drop_on_export(fn):
|
|
attr = get_torchscript_modifier(fn)
|
|
if attr is None:
|
|
return False
|
|
return attr is FunctionModifiers.IGNORE_AND_DROP
|
|
|
|
|
|
def is_ignored_fn(fn):
|
|
mod = get_torchscript_modifier(fn)
|
|
return mod is FunctionModifiers.IGNORE_AND_DROP or mod is FunctionModifiers.IGNORE
|
|
|
|
|
|
def get_torchscript_modifier(fn):
|
|
if not callable(fn):
|
|
return None
|
|
if hasattr(fn, '__func__'):
|
|
fn = fn.__func__
|
|
return getattr(fn, '_torchscript_modifier', FunctionModifiers.DEFAULT)
|
|
|
|
|
|
def _parameter_list(parameter_names_fn):
|
|
"""
|
|
Decorator to denote that a function returns a list of all the parameters
|
|
in a module
|
|
"""
|
|
def decorator(fn):
|
|
fn._parameter_names_fn = parameter_names_fn
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
# overloading registration
|
|
# overloads get registered in this file, and compiled in torch/jit/__init__.py
|
|
# so that they can be imported in nn/functional.py without an import cycle
|
|
|
|
# qualified_name => list[overload_functions]
|
|
_overloaded_fns = {} # noqa: T484
|
|
|
|
def _overload(func):
|
|
qual_name = _qualified_name(func)
|
|
global _overloaded_fns
|
|
fn_overload_list = _overloaded_fns.get(qual_name)
|
|
if fn_overload_list is None:
|
|
fn_overload_list = []
|
|
_overloaded_fns[qual_name] = fn_overload_list
|
|
fn_overload_list.append(func)
|
|
return func
|
|
|
|
def _get_fn_overloads(qual_name):
|
|
return _overloaded_fns.get(qual_name)
|
|
|
|
def _clear_fn_overloads(qual_name):
|
|
del _overloaded_fns[qual_name]
|
|
|
|
try:
|
|
import typing
|
|
from typing import Tuple, List, Dict, Optional
|
|
|
|
def is_tuple(ann):
|
|
# For some reason Python 3.7 violates the Type[A, B].__origin__ == Type rule
|
|
return ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is typing.Tuple or
|
|
getattr(ann, '__origin__', None) is tuple)
|
|
|
|
def is_list(ann):
|
|
return ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is typing.List or
|
|
getattr(ann, '__origin__', None) is list)
|
|
|
|
def is_dict(ann):
|
|
return ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is typing.Dict or
|
|
getattr(ann, '__origin__', None) is dict)
|
|
|
|
def is_optional(ann):
|
|
# Optional[T] is just shorthand for Union[T, None], so check for both
|
|
union_optional = False
|
|
if ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is typing.Union):
|
|
args = getattr(ann, '__args__', ())
|
|
if len(args) == 2:
|
|
union_optional = (issubclass(args[1], type(None)) and not issubclass(args[0], type(None))) \
|
|
or (issubclass(args[0], type(None)) and not issubclass(args[1], type(None)))
|
|
|
|
optional = ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is typing.Optional)
|
|
|
|
return optional or union_optional
|
|
|
|
except ImportError:
|
|
# A minimal polyfill for versions of Python that don't have typing.
|
|
# Note that this means that they also don't support the fancy annotation syntax, so
|
|
# those instances will only be used in our tiny `type: ` comment interpreter.
|
|
|
|
# The __getitem__ in typing is implemented using metaclasses, but I'm too lazy for that.
|
|
class TupleCls(object):
|
|
def __getitem__(self, types):
|
|
return TupleInstance(types)
|
|
|
|
class TupleInstance(object):
|
|
__slots__ = ['__args__']
|
|
|
|
def __init__(self, types):
|
|
self.__args__ = types
|
|
|
|
class ListInstance(object):
|
|
__slots__ = ['__args__']
|
|
|
|
def __init__(self, types):
|
|
self.__args__ = types
|
|
|
|
class ListCls(object):
|
|
def __getitem__(self, types):
|
|
return TupleInstance(types)
|
|
|
|
class DictInstance(object):
|
|
__slots__ = ['__args__']
|
|
|
|
def __init__(self, types):
|
|
self.__args__ = types
|
|
|
|
class DictCls(object):
|
|
def __getitem__(self, types):
|
|
return DictInstance(types)
|
|
|
|
class OptionalInstance(object):
|
|
__slots__ = ['__args__']
|
|
|
|
def __init__(self, types):
|
|
self.__args__ = types
|
|
|
|
class OptionalCls(object):
|
|
def __getitem__(self, types):
|
|
return OptionalInstance(types)
|
|
|
|
Tuple = TupleCls() # noqa: T484
|
|
List = ListCls() # noqa: T484
|
|
Dict = DictCls() # noqa: T484
|
|
Optional = DictCls() # noqa: T484
|
|
|
|
def is_tuple(ann):
|
|
return isinstance(ann, TupleInstance)
|
|
|
|
def is_list(ann):
|
|
return isinstance(ann, ListInstance)
|
|
|
|
def is_dict(ann):
|
|
return isinstance(ann, DictInstance)
|
|
|
|
def is_optional(ann):
|
|
return isinstance(ann, OptionalInstance)
|
|
|
|
|
|
try:
|
|
import typing_extensions
|
|
from typing_extensions import Final
|
|
|
|
def is_final(ann):
|
|
return ann.__module__ == 'typing_extensions' and \
|
|
(getattr(ann, '__origin__', None) is typing_extensions.Final)
|
|
except ImportError:
|
|
# Same as above, this polyfill is only for `typing_extensions`
|
|
class FinalInstance(object):
|
|
__slots__ = ['__args__']
|
|
|
|
def __init__(self, types):
|
|
self.__args__ = types
|
|
|
|
class FinalCls(object):
|
|
def __getitem__(self, types):
|
|
return FinalInstance(types)
|
|
|
|
Final = FinalCls() # noqa: T484
|
|
|
|
def is_final(ann):
|
|
return isinstance(ann, FinalInstance)
|
|
|
|
|
|
# allows BroadcastingList instance to be subscriptable
|
|
class BroadcastingListCls(object):
|
|
def __getitem__(self, types):
|
|
return
|
|
|
|
# mypy doesn't support parameters on types, so we have to explicitly type each
|
|
# list size
|
|
BroadcastingList1 = BroadcastingListCls()
|
|
for i in range(2, 7):
|
|
globals()["BroadcastingList{}".format(i)] = BroadcastingList1
|
|
|
|
# Retrieves a fully-qualified name (module hierarchy + classname) for a given obj.
|
|
def _qualified_name(obj):
|
|
# short-circuit in cases where the object already has a known qualified name
|
|
if isinstance(obj, torch._C.Function):
|
|
return obj.qualified_name
|
|
|
|
name = obj.__name__
|
|
module_name = obj.__module__
|
|
|
|
# If the module is actually a torchbind module, then we should short circuit
|
|
if module_name == "torch._classes":
|
|
return obj.qualified_name
|
|
|
|
# The Python docs are very clear that `__module__` can be None, but I can't
|
|
# figure out when it actually would be.
|
|
if module_name is None:
|
|
raise RuntimeError("Could not get qualified name for class '{}': "
|
|
"__module__ can't be None.".format(name))
|
|
|
|
# if getattr(sys.modules[module_name], name) is not obj:
|
|
# raise RuntimeError("Could not get qualified name for class '{}': "
|
|
# "the attr {} on module {} is not the the class".format(name, name, module_name))
|
|
|
|
# __main__ is a builtin module, so rewrite it to "__torch__".
|
|
if module_name == "__main__":
|
|
module_name = "__torch__"
|
|
else:
|
|
# Everything else gets a "__torch__" prefix to avoid name collisions
|
|
# with the names of user values.
|
|
module_name = "__torch__." + module_name
|
|
|
|
if "." in name:
|
|
raise RuntimeError("Could not get qualified name for class '{}': "
|
|
"'{}' is not a valid identifier".format(name, name))
|
|
|
|
return module_name + "." + name
|