mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
b0ac9bfb2b
Summary: This adds a comment above `should_drop` to prevent someone from inadvertently breaking JIT coverage by renaming the function without updating the correct references. The current JIT plug-in uses `should_drop` to figure out which code is going to be JIT'd. If the function is named differently, the plug-in would also need to be updated. Question: I understand this may not be the cleanest solution. Would a cleaner solution be to create a dummy function that would simply exist for the JIT plug-in? I did not immediately do that as that may be adding unnecessary code complexity in torch.jit. Pull Request resolved: https://github.com/pytorch/pytorch/pull/57961 Reviewed By: samestep Differential Revision: D28933587 Pulled By: janeyx99 fbshipit-source-id: 260aaf7b11f07de84a81d6c3554c4a5ce479d623
1172 lines
42 KiB
Python
1172 lines
42 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 contextlib
|
|
import collections
|
|
import enum
|
|
import inspect
|
|
import ast
|
|
import weakref
|
|
import warnings
|
|
from textwrap import dedent
|
|
import torch
|
|
import sys
|
|
import builtins
|
|
import io
|
|
import pickle
|
|
import functools
|
|
# This is needed. `torch._jit_internal` is imported before `torch.distributed.__init__`.
|
|
# Explicitly ask to import `torch.distributed.__init__` first.
|
|
# Otherwise, "AttributeError: module 'torch' has no attribute 'distributed'" is raised.
|
|
import torch.distributed.rpc
|
|
from torch._utils_internal import get_source_lines_and_file
|
|
from torch.futures import Future
|
|
import torch.package._mangling as package_mangling
|
|
from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, TypeVar, Union # noqa: F401
|
|
|
|
if sys.version_info[:2] > (3, 7):
|
|
from typing import Final
|
|
else:
|
|
from typing_extensions import Final
|
|
|
|
# Wrapper functions that can call either of 2 functions depending on a boolean
|
|
# argument
|
|
boolean_dispatched: 'weakref.WeakKeyDictionary[Callable, Dict[str, Callable]]' = weakref.WeakKeyDictionary() # noqa: T484
|
|
|
|
|
|
def createResolutionCallbackFromEnv(lookup_base):
|
|
"""
|
|
Creates a resolution callback that will look up qualified names in an
|
|
environment, starting with `lookup_base` for the base of any qualified
|
|
names, then proceeding down the lookup chain with the resolved object.
|
|
|
|
You should not use this directly, it should only be used from the other
|
|
createResolutionCallbackFrom* functions.
|
|
"""
|
|
def lookupInModule(qualified_name, module):
|
|
if '.' in qualified_name:
|
|
parts = qualified_name.split('.')
|
|
base = parts[0]
|
|
remaining_pieces = '.'.join(parts[1:])
|
|
module_value = getattr(module, base)
|
|
return lookupInModule(remaining_pieces, module_value)
|
|
else:
|
|
return getattr(module, qualified_name)
|
|
|
|
def parseNestedExpr(expr, module) -> Tuple[Any, int]:
|
|
i = 0
|
|
while i < len(expr) and expr[i] not in (',', '[', ']'):
|
|
i += 1
|
|
|
|
# Special case logic for the empty Tuple as a subscript (used
|
|
# in the type annotation `Tuple[()]`)
|
|
if expr[:i] == '()':
|
|
return (), i
|
|
|
|
base = lookupInModule(expr[:i].strip(), module)
|
|
assert base is not None, f"Unresolvable type {expr[:i]}"
|
|
if i == len(expr) or expr[i] != '[':
|
|
return base, i
|
|
|
|
assert expr[i] == '['
|
|
parts = []
|
|
while expr[i] != ']':
|
|
part_len = 0
|
|
i += 1
|
|
part, part_len = parseNestedExpr(expr[i:], module)
|
|
parts.append(part)
|
|
i += part_len
|
|
if len(parts) > 1:
|
|
return base[tuple(parts)], i + 1
|
|
else:
|
|
return base[parts[0]], i + 1
|
|
|
|
def parseExpr(expr, module):
|
|
try:
|
|
value, len_parsed = parseNestedExpr(expr, module)
|
|
assert len_parsed == len(expr), "whole expression was not parsed, falling back to c++ parser"
|
|
return value
|
|
except Exception:
|
|
"""
|
|
The python resolver fails in several cases in known unit tests, and is intended
|
|
to fall back gracefully to the c++ resolver in general. For example, python 2 style
|
|
annotations which are frequent in our unit tests often fail with types e.g. int not
|
|
resolvable from the calling frame.
|
|
"""
|
|
return None
|
|
|
|
return lambda expr: parseExpr(expr, lookup_base)
|
|
|
|
|
|
def createResolutionCallbackFromFrame(frames_up: int = 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 createResolutionCallbackFromFrame (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 createResolutionCallbackFromFrame. Also for example, if frames_up is set
|
|
to 1, then the frame of the caller's caller of createResolutionCallbackFromFrame
|
|
will be taken.
|
|
|
|
For example, the following program prints 2::
|
|
|
|
def bar():
|
|
cb = createResolutionCallbackFromFrame(1)
|
|
print(cb("foo"))
|
|
|
|
def baz():
|
|
foo = 2
|
|
bar()
|
|
|
|
baz()
|
|
"""
|
|
frame = inspect.currentframe()
|
|
i = 0
|
|
while i < frames_up + 1:
|
|
assert frame is not None
|
|
frame = frame.f_back
|
|
i += 1
|
|
|
|
assert frame is not None
|
|
f_locals = frame.f_locals
|
|
f_globals = frame.f_globals
|
|
|
|
class env(object):
|
|
def __getattr__(self, key):
|
|
if key in f_locals:
|
|
return f_locals[key]
|
|
elif key in f_globals:
|
|
return f_globals[key]
|
|
elif key in dir(builtins):
|
|
return getattr(builtins, key)
|
|
|
|
return createResolutionCallbackFromEnv(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
|
|
|
|
# [local resolution in python]
|
|
# Depending on where a variable is defined, and where it is used, we may
|
|
# or may not be able to recover its value when recursively compiling a
|
|
# script function. Remember in the general case, a module or function is
|
|
# first defined and then later scripted. This means we do not have a
|
|
# chance to capture the active frames when the function is defined. Hence any
|
|
# name resolution has to happen later on the created closure. The way
|
|
# python captures type annotations restricts what we can recover. The
|
|
# follow example illustrates the different cases:
|
|
#
|
|
# class MyGlobalClass:
|
|
# ...
|
|
# def my_local_scope():
|
|
# @torch.jit.script
|
|
# class MyClass:
|
|
# ...
|
|
# @torch.jit.script
|
|
# class MyClassUsedAsVar:
|
|
# ...
|
|
# def eg(x: MyClass, y: MyGlobalClass):
|
|
# a_local_capture : Foo
|
|
# return MyClassUsedAsVar(x)
|
|
#
|
|
# MyGlobalClass is defined in the __globals__ dictionary of function
|
|
# 'eg', so it is always recoverable. my_local_scope introduces a new local
|
|
# variable scope in the function. Classes defined here are only visible as
|
|
# local variables. For the case of MyClassUsedAsVar, it is captured
|
|
# because it is used as a variable inside the body of the function, and we
|
|
# can resolve it using the captures returned from `get_closure`. However,
|
|
# the type annotations are not captured by the closure. In Python
|
|
# 3.0--3.9, the _value_ of MyClass and MyGlobalClass will be available as
|
|
# annotations on `eg``, but starting in Python 4.0, they will represented as
|
|
# strings and no longer present. Furthermore, since the body of `eg` does
|
|
# not reference those names, they do not appear in the list of closed over
|
|
# variables. In Python 2.x, type annotations are in comments, leading to a
|
|
# similar situation where their definitions are not available. We anticipate
|
|
# that most users will not run into this issue because their modules and
|
|
# functions will be defined at a global scope like MyGlobalClass. In cases
|
|
# where they are not, it is possible to work around issues by declaring the
|
|
# values global in the function.
|
|
# In Python 3.9 declaring class as global will make it invisible to
|
|
# `inspect.getsource`, see https://bugs.python.org/issue42666 .
|
|
# This could be worked around by manualy adding it to `global()` dictionary.
|
|
|
|
|
|
|
|
def createResolutionCallbackFromClosure(fn):
|
|
"""
|
|
Create a resolutionCallback by introspecting the function instead of
|
|
looking up the stack for the enclosing scope
|
|
"""
|
|
closure = get_closure(fn)
|
|
|
|
class closure_lookup(object):
|
|
# This is a class since `closure` is a dict and it's easier in
|
|
# `env_helper` if everything just works with `getattr` calls
|
|
def __getattr__(self, key):
|
|
if key in closure:
|
|
return closure[key]
|
|
elif hasattr(builtins, key):
|
|
return getattr(builtins, key)
|
|
return None
|
|
|
|
return createResolutionCallbackFromEnv(closure_lookup())
|
|
|
|
|
|
def can_compile_class(cls) -> bool:
|
|
# 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
|
|
if is_ignored_fn(cls):
|
|
return False
|
|
|
|
# Ignore the following list of built-in classes.
|
|
ignored_builtin_classes = (torch.nn.Module, tuple, list, Exception)
|
|
if issubclass(cls, ignored_builtin_classes):
|
|
return False
|
|
|
|
names = cls.__dict__
|
|
fns = [getattr(cls, name) for name in names if inspect.isroutine(getattr(cls, name, None))]
|
|
has_code = [hasattr(fn, '__code__') for fn in fns]
|
|
return all(has_code)
|
|
|
|
|
|
def get_callable_argument_names(fn) -> List[str]:
|
|
"""
|
|
Gets names of all POSITIONAL_OR_KEYWORD arguments for callable `fn`.
|
|
Returns an empty list when other types of arguments are present.
|
|
|
|
This is used by `torch.jit.trace` to assign meaningful argument names to
|
|
traced functions and modules.
|
|
|
|
Args:
|
|
fn: A callable.
|
|
Returns:
|
|
Argument names: List[str]
|
|
"""
|
|
# inspect.signature may fail, give up in that case.
|
|
try:
|
|
callable_signature = inspect.signature(fn)
|
|
except Exception:
|
|
return []
|
|
|
|
argument_names = []
|
|
for name, param in callable_signature.parameters.items():
|
|
# All four other types of arguments do not map to individual values
|
|
# with a keyword as name.
|
|
if not param.kind == param.POSITIONAL_OR_KEYWORD:
|
|
return []
|
|
|
|
argument_names.append(name)
|
|
|
|
return argument_names
|
|
|
|
|
|
def get_annotation_str(annotation):
|
|
"""
|
|
Convert an AST node containing a type annotation to the string present in the source
|
|
that represents the same annotation.
|
|
"""
|
|
if isinstance(annotation, ast.Name):
|
|
return annotation.id
|
|
elif isinstance(annotation, ast.Attribute):
|
|
return '.'.join([get_annotation_str(annotation.value), annotation.attr])
|
|
elif isinstance(annotation, ast.Subscript):
|
|
# In Python3.9+ subscript indicies are not wrapped in ast.Index
|
|
subscript_slice = annotation.slice if sys.version_info >= (3, 9) else annotation.slice.value # type: ignore[attr-defined]
|
|
return f"{get_annotation_str(annotation.value)}[{get_annotation_str(subscript_slice)}]"
|
|
elif isinstance(annotation, ast.Tuple):
|
|
return ','.join([get_annotation_str(elt) for elt in annotation.elts])
|
|
elif isinstance(annotation, ast.Constant) or isinstance(annotation, ast.NameConstant):
|
|
return f"{annotation.value}"
|
|
|
|
# If an AST node is not handled here, it's probably handled in ScriptTypeParser.
|
|
return None
|
|
|
|
|
|
def get_type_hint_captures(fn):
|
|
"""
|
|
Get a dictionary containing type resolution mappings necessary to resolve types
|
|
for the literal annotations on 'fn'. These are not considered to be closed-over by fn
|
|
and must be obtained separately (e.g. using this function).
|
|
|
|
Args:
|
|
fn: A callable.
|
|
Returns:
|
|
A Dict[str, Any] containing a mapping from the literal annotations used on
|
|
fn to the Python objects they refer to.
|
|
"""
|
|
# Gather a dictionary of parameter name -> type, skipping any parameters whose annotated
|
|
# types are strings. These are only understood by TorchScript in the context of a type annotation
|
|
# that refers to a class in its own definition, but trying to include a mapping for this in the result
|
|
# function would cause infinite recursion because the class is currently being compiled.
|
|
# In addition, there is logic in ScriptTypeParser to handle this.
|
|
signature = inspect.signature(fn)
|
|
name_to_type = {
|
|
name: parameter.annotation
|
|
for name, parameter in signature.parameters.items()
|
|
if parameter.annotation is not inspect.Parameter.empty and not isinstance(parameter.annotation, str)
|
|
}
|
|
|
|
# Then, get the literal type annotations from the function declaration
|
|
# by source inspection. This accounts for the case in which aliases are used
|
|
# to annotate the arguments (e.g device_t = torch.device, and then d: device_t).
|
|
src = inspect.getsource(fn)
|
|
|
|
# frontend.py cannot be used here because it includes _jit_internal, so use ast instead.
|
|
a = ast.parse(dedent(src))
|
|
if len(a.body) != 1 or not isinstance(a.body[0], ast.FunctionDef):
|
|
raise RuntimeError(f"Expected {fn} to be a function")
|
|
f = a.body[0]
|
|
|
|
# Prepare a dictionary of source annotation -> type, which will be the final result of this function,
|
|
# by using the parsed AST (f) to reconstruct source annotations as strings for each parameter and mapping
|
|
# them to the type object corresponding to the annotation via name_to_type using the parameter name.
|
|
annotation_to_type = {}
|
|
|
|
for arg in f.args.args:
|
|
# Get the source type annotation string for this argument if possible.
|
|
arg_annotation_str = get_annotation_str(arg.annotation) if arg.annotation else None
|
|
|
|
# If the argument has no annotation or get_annotation_str cannot convert it to a string,
|
|
# arg_annotation_str will be None. Skip this arg; ScriptTypeParser will probably handle
|
|
# this in the latter case.
|
|
if arg_annotation_str is None:
|
|
continue
|
|
|
|
# Insert {arg_annotation_str: type} into annotation_to_type if possible. One reason arg_name may not
|
|
# be present in name_to_type is that the annotation itself is a string and not a type object
|
|
# (common for self-refential annotations in classes). Once again, let ScriptTypeParser handle this.
|
|
arg_name = arg.arg
|
|
if arg_name in name_to_type:
|
|
annotation_to_type[arg_annotation_str] = name_to_type[arg_name]
|
|
|
|
# If there is a valid return annotation, include it in annotation_to_type. As with argument annotations,
|
|
# the literal annotation has to be convertible to a string by get_annotation_str, and the actual type
|
|
# of the annotation cannot be a string.
|
|
literal_return_annotation = get_annotation_str(f.returns)
|
|
valid_literal_annotation = literal_return_annotation is not None
|
|
return_annotation = signature.return_annotation
|
|
valid_return_annotation_type = return_annotation is not inspect.Parameter.empty and not isinstance(return_annotation, str)
|
|
if valid_literal_annotation and valid_return_annotation_type:
|
|
annotation_to_type[literal_return_annotation] = return_annotation
|
|
|
|
return annotation_to_type
|
|
|
|
|
|
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))
|
|
captures.update(get_type_hint_captures(fn))
|
|
|
|
def lookup_in_class(key):
|
|
if key in captures:
|
|
return captures[key]
|
|
else:
|
|
return getattr(builtins, key, None)
|
|
|
|
return lookup_in_class
|
|
|
|
|
|
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.
|
|
"""
|
|
UNUSED = "unused (ignored and replaced with raising of an exception)"
|
|
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)"
|
|
COPY_TO_SCRIPT_WRAPPER = \
|
|
"if this method is not scripted, copy the python method onto the scripted model"
|
|
|
|
|
|
def export(fn):
|
|
"""
|
|
This decorator indicates that a method on an ``nn.Module`` is used as an entry point into a
|
|
:class:`ScriptModule` and should be compiled.
|
|
|
|
``forward`` implicitly is assumed to be an entry point, so it does not need this decorator.
|
|
Functions and methods called from ``forward`` are compiled as they are seen
|
|
by the compiler, so they do not need this decorator either.
|
|
|
|
Example (using ``@torch.jit.export`` on a method):
|
|
|
|
.. testcode::
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class MyModule(nn.Module):
|
|
def implicitly_compiled_method(self, x):
|
|
return x + 99
|
|
|
|
# `forward` is implicitly decorated with `@torch.jit.export`,
|
|
# so adding it here would have no effect
|
|
def forward(self, x):
|
|
return x + 10
|
|
|
|
@torch.jit.export
|
|
def another_forward(self, x):
|
|
# When the compiler sees this call, it will compile
|
|
# `implicitly_compiled_method`
|
|
return self.implicitly_compiled_method(x)
|
|
|
|
def unused_method(self, x):
|
|
return x - 20
|
|
|
|
# `m` will contain compiled methods:
|
|
# `forward`
|
|
# `another_forward`
|
|
# `implicitly_compiled_method`
|
|
# `unused_method` will not be compiled since it was not called from
|
|
# any compiled methods and wasn't decorated with `@torch.jit.export`
|
|
m = torch.jit.script(MyModule())
|
|
"""
|
|
fn._torchscript_modifier = FunctionModifiers.EXPORT
|
|
return fn
|
|
|
|
|
|
def unused(fn):
|
|
"""
|
|
This decorator indicates to the compiler that a function or method should
|
|
be ignored and replaced with the raising of an exception. This allows you
|
|
to leave code in your model that is not yet TorchScript compatible and still
|
|
export your model.
|
|
|
|
Example (using ``@torch.jit.unused`` on a method)::
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class MyModule(nn.Module):
|
|
def __init__(self, use_memory_efficient):
|
|
super(MyModule, self).__init__()
|
|
self.use_memory_efficient = use_memory_efficient
|
|
|
|
@torch.jit.unused
|
|
def memory_efficient(self, x):
|
|
import pdb
|
|
pdb.set_trace()
|
|
return x + 10
|
|
|
|
def forward(self, x):
|
|
# Use not-yet-scriptable memory efficient mode
|
|
if self.use_memory_efficient:
|
|
return self.memory_efficient(x)
|
|
else:
|
|
return x + 10
|
|
|
|
m = torch.jit.script(MyModule(use_memory_efficient=False))
|
|
m.save("m.pt")
|
|
|
|
m = torch.jit.script(MyModule(use_memory_efficient=True))
|
|
# exception raised
|
|
m(torch.rand(100))
|
|
"""
|
|
if isinstance(fn, property):
|
|
prop = fn
|
|
setattr(prop.fget, "_torchscript_modifier", FunctionModifiers.UNUSED) # noqa: B010
|
|
|
|
if prop.fset:
|
|
setattr(prop.fset, "_torchscript_modifier", FunctionModifiers.UNUSED) # noqa: B010
|
|
|
|
return prop
|
|
|
|
fn._torchscript_modifier = FunctionModifiers.UNUSED
|
|
return fn
|
|
|
|
# No op context manager from python side
|
|
class _IgnoreContextManager(contextlib.AbstractContextManager):
|
|
def __init__(self, **kwargs):
|
|
pass
|
|
|
|
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
|
pass
|
|
|
|
def ignore(drop=False, **kwargs):
|
|
"""
|
|
This decorator indicates to the compiler that a function or method should
|
|
be ignored and left as a Python function. This allows you to leave code in
|
|
your model that is not yet TorchScript compatible. If called from TorchScript,
|
|
ignored functions will dispatch the call to the Python interpreter. Models with ignored
|
|
functions cannot be exported; use :func:`@torch.jit.unused <torch.jit.unused>` instead.
|
|
|
|
Example (using ``@torch.jit.ignore`` on a method)::
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class MyModule(nn.Module):
|
|
@torch.jit.ignore
|
|
def debugger(self, x):
|
|
import pdb
|
|
pdb.set_trace()
|
|
|
|
def forward(self, x):
|
|
x += 10
|
|
# The compiler would normally try to compile `debugger`,
|
|
# but since it is `@ignore`d, it will be left as a call
|
|
# to Python
|
|
self.debugger(x)
|
|
return x
|
|
|
|
m = torch.jit.script(MyModule())
|
|
|
|
# Error! The call `debugger` cannot be saved since it calls into Python
|
|
m.save("m.pt")
|
|
|
|
Example (using ``@torch.jit.ignore(drop=True)`` on a method):
|
|
|
|
.. testcode::
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class MyModule(nn.Module):
|
|
@torch.jit.ignore(drop=True)
|
|
def training_method(self, x):
|
|
import pdb
|
|
pdb.set_trace()
|
|
|
|
def forward(self, x):
|
|
if self.training:
|
|
self.training_method(x)
|
|
return x
|
|
|
|
m = torch.jit.script(MyModule())
|
|
|
|
# This is OK since `training_method` is not saved, the call is replaced
|
|
# with a `raise`.
|
|
m.save("m.pt")
|
|
|
|
.. testcleanup::
|
|
|
|
import os
|
|
os.remove('m.pt')
|
|
"""
|
|
|
|
if callable(drop):
|
|
# used without any args, so drop is actually a function
|
|
# @torch.jit.ignore
|
|
# def fn(...):
|
|
fn = drop
|
|
fn._torchscript_modifier = FunctionModifiers.IGNORE
|
|
return fn
|
|
|
|
if not isinstance(drop, bool):
|
|
raise RuntimeError("Argument to @torch.jit.ignore must be a bool or "
|
|
f"a function but got {drop}")
|
|
|
|
# for backwards compat
|
|
drop_on_export = kwargs.pop("drop_on_export", None)
|
|
if drop_on_export:
|
|
warnings.warn("ignore(drop_on_export=True) has been deprecated. TorchScript will now drop the function "
|
|
"call on compilation. Use torch.jit.unused now. {}", category=FutureWarning)
|
|
|
|
drop = drop_on_export
|
|
elif drop:
|
|
warnings.warn("ignore(True) has been deprecated. TorchScript will now drop the function "
|
|
"call on compilation. Use torch.jit.unused now. {}", category=FutureWarning)
|
|
|
|
def decorator(fn):
|
|
if drop:
|
|
fn._torchscript_modifier = FunctionModifiers.UNUSED
|
|
else:
|
|
fn._torchscript_modifier = FunctionModifiers.IGNORE
|
|
return fn
|
|
return decorator
|
|
|
|
|
|
def _copy_to_script_wrapper(fn):
|
|
fn._torchscript_modifier = FunctionModifiers.COPY_TO_SCRIPT_WRAPPER
|
|
return fn
|
|
|
|
def module_has_exports(mod):
|
|
for name in dir(mod):
|
|
if hasattr(mod, name):
|
|
item = getattr(mod, name)
|
|
if callable(item):
|
|
if get_torchscript_modifier(item) is FunctionModifiers.EXPORT:
|
|
return True
|
|
return False
|
|
|
|
|
|
# WARNING: should_drop is currently being used by our JIT code coverage plug-in to mark JIT'd code as covered. If you
|
|
# rename this function, please update references in tools/coverage_plugins_package/src/coverage_plugins/jit_plugin.py to
|
|
# allow JIT'd code to still be covered.
|
|
def should_drop(fn) -> bool:
|
|
attr = get_torchscript_modifier(fn)
|
|
if attr is None:
|
|
return False
|
|
return attr is FunctionModifiers.UNUSED
|
|
|
|
|
|
def is_ignored_fn(fn) -> bool:
|
|
mod = get_torchscript_modifier(fn)
|
|
return mod is FunctionModifiers.UNUSED or mod is FunctionModifiers.IGNORE
|
|
|
|
|
|
def is_static_fn(cls, fn) -> bool:
|
|
return isinstance(inspect.getattr_static(cls, fn, default=None), staticmethod)
|
|
|
|
def get_static_fn(cls, fn):
|
|
return inspect.getattr_static(cls, fn).__func__
|
|
|
|
|
|
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 copy_torchscript_modifier(orig, new) -> None:
|
|
attr = get_torchscript_modifier(orig)
|
|
if attr is None:
|
|
return
|
|
new._torchscript_modifier = attr
|
|
|
|
# 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 : Dict[str, List[Callable]] = {} # 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) -> None:
|
|
del _overloaded_fns[qual_name]
|
|
|
|
def get_class_name_lineno(method) -> Tuple[str, int]:
|
|
current_frame = inspect.currentframe()
|
|
|
|
# one for the get_class_name call, one for _overload_method call
|
|
for i in range(2):
|
|
assert current_frame is not None # assert current frame is not an Optional[FrameType]
|
|
current_frame = current_frame.f_back
|
|
|
|
assert current_frame is not None # same here
|
|
class_name = current_frame.f_code.co_name
|
|
line_no = current_frame.f_code.co_firstlineno
|
|
return class_name, line_no
|
|
|
|
# At the the point the decorator is applied to class methods the method
|
|
# has no reference to its owning class. _qualified_name would not include
|
|
# the class it is defined in, so any methods with the same name in the same file
|
|
# would have the same _qualified_name, even if they were defined in different
|
|
# classes. This problem only exists in python 2.
|
|
# We get around this problem by looking at the stack frame and identifying
|
|
# the class name, and throwing an error whenever overloads are used
|
|
# when modules of the same name are in the same file
|
|
|
|
# qualified_name => class name => list[overload_functions]
|
|
_overloaded_methods : Dict[str, Dict[str, List[Callable]]] = {} # noqa: T484
|
|
|
|
|
|
# (qualified_name, class name) => class_fileno
|
|
_overloaded_method_class_fileno = {}
|
|
|
|
def _overload_method(func):
|
|
qual_name = _qualified_name(func)
|
|
global _overloaded_methods
|
|
class_name_map = _overloaded_methods.get(qual_name, None)
|
|
if class_name_map is None:
|
|
class_name_map = {}
|
|
_overloaded_methods[qual_name] = class_name_map
|
|
|
|
class_name, line_no = get_class_name_lineno(func)
|
|
method_overloads = class_name_map.get(class_name, None)
|
|
if method_overloads is None:
|
|
method_overloads = []
|
|
class_name_map[class_name] = method_overloads
|
|
_overloaded_method_class_fileno[(qual_name, class_name)] = line_no
|
|
else:
|
|
existing_lineno = _overloaded_method_class_fileno[(qual_name, class_name)]
|
|
if existing_lineno != line_no:
|
|
raise RuntimeError("Cannot currently overload the same method name in two different"
|
|
" classes with the same name in the same module")
|
|
|
|
method_overloads.append(func)
|
|
return func
|
|
|
|
def _get_overloaded_methods(method, mod_class):
|
|
# TODO: __name__ not set for submodules in recursive script
|
|
if not hasattr(method, "__name__"):
|
|
return None
|
|
qual_name = _qualified_name(method)
|
|
class_name_map = _overloaded_methods.get(qual_name, None)
|
|
if class_name_map is None:
|
|
return None
|
|
overloads = class_name_map.get(mod_class.__name__, None)
|
|
if overloads is None:
|
|
return None
|
|
|
|
method_line_no = get_source_lines_and_file(method)[1]
|
|
mod_class_fileno = get_source_lines_and_file(mod_class)[1]
|
|
mod_end_fileno = mod_class_fileno + len(get_source_lines_and_file(mod_class)[0])
|
|
if not (method_line_no >= mod_class_fileno and method_line_no <= mod_end_fileno):
|
|
raise Exception("Overloads are not useable when a module is redeclared within the same file: " + str(method))
|
|
return overloads
|
|
|
|
|
|
def is_tuple(ann) -> bool:
|
|
if ann is Tuple:
|
|
raise_error_container_parameter_missing("Tuple")
|
|
|
|
# For some reason Python 3.7 violates the Type[A, B].__origin__ == Type rule
|
|
if not hasattr(ann, '__module__'):
|
|
return False
|
|
return ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is Tuple or
|
|
getattr(ann, '__origin__', None) is tuple)
|
|
|
|
def is_list(ann) -> bool:
|
|
if ann is List:
|
|
raise_error_container_parameter_missing("List")
|
|
|
|
if not hasattr(ann, '__module__'):
|
|
return False
|
|
return ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is List or
|
|
getattr(ann, '__origin__', None) is list)
|
|
|
|
def is_dict(ann) -> bool:
|
|
if ann is Dict:
|
|
raise_error_container_parameter_missing("Dict")
|
|
|
|
if not hasattr(ann, '__module__'):
|
|
return False
|
|
return ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is Dict or
|
|
getattr(ann, '__origin__', None) is dict)
|
|
|
|
def is_optional(ann) -> bool:
|
|
if ann is Optional:
|
|
raise_error_container_parameter_missing("Optional")
|
|
|
|
# Optional[T] is just shorthand for Union[T, None], so check for both
|
|
def safe_is_subclass(the_type, super_type):
|
|
# Don't throw if `the_type` isn't a class type (e.g. if it is
|
|
# another type annotation instance)
|
|
if not inspect.isclass(the_type):
|
|
return False
|
|
return issubclass(the_type, super_type)
|
|
|
|
if not hasattr(ann, '__module__'):
|
|
return False
|
|
|
|
union_optional = False
|
|
if ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is Union):
|
|
args = getattr(ann, '__args__', ())
|
|
if len(args) == 2:
|
|
union_optional = (safe_is_subclass(args[1], type(None)) and not safe_is_subclass(args[0], type(None))) \
|
|
or (safe_is_subclass(args[0], type(None)) and not safe_is_subclass(args[1], type(None)))
|
|
|
|
optional = ann.__module__ == 'typing' and \
|
|
(getattr(ann, '__origin__', None) is Optional)
|
|
|
|
return optional or union_optional
|
|
|
|
def is_future(ann) -> bool:
|
|
if ann is Future:
|
|
raise RuntimeError(
|
|
"Attempted to use Future without a "
|
|
"contained type. Please add a contained type, e.g. "
|
|
"Future[int]"
|
|
)
|
|
return getattr(ann, "__origin__", None) is Future
|
|
|
|
if torch.distributed.rpc.is_available():
|
|
from torch.distributed.rpc import RRef
|
|
|
|
def is_rref(ann) -> bool:
|
|
if ann is RRef:
|
|
raise RuntimeError(
|
|
"Attempted to use RRef without a "
|
|
"contained type. Please add a contained type, e.g. "
|
|
"RRef[int]"
|
|
)
|
|
return getattr(ann, "__origin__", None) is RRef
|
|
|
|
def is_final(ann) -> bool:
|
|
return ann.__module__ in {'typing', 'typing_extensions'} and \
|
|
(getattr(ann, '__origin__', None) is Final or isinstance(ann, type(Final)))
|
|
|
|
# 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()[f"BroadcastingList{i}"] = BroadcastingList1
|
|
|
|
|
|
def is_scripting() -> bool:
|
|
r"""
|
|
Function that returns True when in compilation and False otherwise. This
|
|
is useful especially with the @unused decorator to leave code in your
|
|
model that is not yet TorchScript compatible.
|
|
.. testcode::
|
|
|
|
import torch
|
|
|
|
@torch.jit.unused
|
|
def unsupported_linear_op(x):
|
|
return x
|
|
|
|
def linear(x):
|
|
if torch.jit.is_scripting():
|
|
return torch.linear(x)
|
|
else:
|
|
return unsupported_linear_op(x)
|
|
"""
|
|
return False
|
|
|
|
|
|
# Retrieves a fully-qualified name (module hierarchy + classname) for a given obj.
|
|
def _qualified_name(obj) -> str:
|
|
# This special case allows us to override the qualified name on a type.
|
|
# It's currently used in conjunction with tracing, where we create a
|
|
# fake module to filter only supported attributes. However, since this
|
|
# new type is defined as a local class, we need a mechanism to override
|
|
# its qualname so it appears correctly in the TorchScript system. This,
|
|
# we set '_jit_override_qualname' with the original traced module's
|
|
# qualified name, which is picked up here
|
|
if hasattr(obj, '_jit_override_qualname'):
|
|
return obj._jit_override_qualname
|
|
# short-circuit in cases where the object already has a known qualified name
|
|
if isinstance(obj, torch._C.ScriptFunction):
|
|
return obj.qualified_name
|
|
|
|
if getattr(obj, "__name__", None):
|
|
name = obj.__name__
|
|
# Enum classes do not have `__name__` attr, instead they have `name`.
|
|
elif isinstance(obj, enum.Enum):
|
|
name = obj.name
|
|
else:
|
|
raise RuntimeError("Could not get name of python class object")
|
|
|
|
|
|
if name == '<lambda>':
|
|
name = '_lambda' # make name a valid identifier
|
|
|
|
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(f"Could not get qualified name for class '{name}': "
|
|
"__module__ can't be None.")
|
|
|
|
# if getattr(sys.modules[module_name], name) is not obj:
|
|
# raise RuntimeError(f"Could not get qualified name for class '{name}': "
|
|
# f"the attr {name} on module {module_name} is not the the class")
|
|
|
|
# torch.package and TorchScript have separate mangling schemes to avoid
|
|
# name collisions from multiple packages. To avoid them interfering with
|
|
# each other, remove the package mangling here.
|
|
module_name = package_mangling.demangle(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(f"Could not get qualified name for class '{name}': "
|
|
f"'{name}' is not a valid identifier")
|
|
|
|
return module_name + "." + name
|
|
|
|
|
|
# Thin wrapper around SourceRangeFactory to store extra metadata
|
|
# about the function-to-be-compiled.
|
|
class SourceContext(torch._C._jit_tree_views.SourceRangeFactory):
|
|
def __init__(self, source, filename, file_lineno, leading_whitespace_len, uses_true_division=True):
|
|
super(SourceContext, self).__init__(source, filename, file_lineno, leading_whitespace_len)
|
|
self.uses_true_division = uses_true_division
|
|
self.filename = filename
|
|
|
|
@functools.lru_cache(maxsize=None)
|
|
def make_source_context(*args):
|
|
return SourceContext(*args)
|
|
|
|
def fake_range():
|
|
return SourceContext('', None, 0, 0).make_raw_range(0, 1)
|
|
|
|
|
|
def _try_get_dispatched_fn(fn):
|
|
if not callable(fn):
|
|
return None
|
|
return boolean_dispatched.get(fn)
|
|
|
|
|
|
def _get_named_tuple_properties(obj):
|
|
assert issubclass(obj, tuple) and hasattr(obj, '_fields')
|
|
fields = list(obj._fields)
|
|
annotations = []
|
|
has_annotations = hasattr(obj, '__annotations__')
|
|
for field in fields:
|
|
if has_annotations and field in obj.__annotations__:
|
|
the_type = torch.jit.annotations.ann_to_type(obj.__annotations__[field], fake_range())
|
|
annotations.append(the_type)
|
|
else:
|
|
annotations.append(torch._C.TensorType.getInferred())
|
|
return type(obj).__name__, fields, annotations
|
|
|
|
|
|
def _create_named_tuple(t, unqual_name: str, field_names: List[str]):
|
|
# mypy: namedtuple() expects a string literal as the first argument
|
|
TupleType = collections.namedtuple(unqual_name, field_names) # type: ignore[misc]
|
|
return TupleType(*t)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _disable_emit_hooks():
|
|
hooks = torch._C._jit_get_emit_hooks()
|
|
torch._C._jit_set_emit_hooks(None, None)
|
|
yield
|
|
torch._C._jit_set_emit_hooks(hooks[0], hooks[1])
|
|
|
|
|
|
def _disable_emit_hooks_decorator(_DecoratorContextManager) -> None: # noqa: F811
|
|
def __enter__(self) -> None:
|
|
self.hooks = torch._C._jit_get_emit_hooks()
|
|
torch._C._jit_set_emit_hooks(None, None)
|
|
|
|
def __exit__(self, *args) -> None:
|
|
torch._C._jit_set_emit_hooks(self.hooks[0], self.hooks[1])
|
|
|
|
def _is_exception(obj) -> bool:
|
|
if not inspect.isclass(obj):
|
|
return False
|
|
return issubclass(obj, Exception)
|
|
|
|
def raise_error_container_parameter_missing(target_type) -> None:
|
|
if target_type == 'Dict':
|
|
raise RuntimeError(
|
|
"Attempted to use Dict without "
|
|
"contained types. Please add contained type, e.g. "
|
|
"Dict[int, int]"
|
|
)
|
|
raise RuntimeError(
|
|
f"Attempted to use {target_type} without a "
|
|
"contained type. Please add a contained type, e.g. "
|
|
f"{target_type}[int]"
|
|
)
|
|
|
|
|
|
def get_origin(target_type):
|
|
return getattr(target_type, "__origin__", None)
|
|
|
|
|
|
def get_args(target_type):
|
|
return getattr(target_type, "__args__", None)
|
|
|
|
|
|
def check_args_exist(target_type) -> None:
|
|
if target_type is List or target_type is list:
|
|
raise_error_container_parameter_missing("List")
|
|
elif target_type is Tuple or target_type is tuple:
|
|
raise_error_container_parameter_missing("Tuple")
|
|
elif target_type is Dict or target_type is dict:
|
|
raise_error_container_parameter_missing("Dict")
|
|
elif target_type is None or target_type is Optional:
|
|
raise_error_container_parameter_missing("Optional")
|
|
|
|
|
|
# supports List/Dict/Tuple and Optional types
|
|
# TODO support future
|
|
def container_checker(obj, target_type) -> bool:
|
|
origin_type = get_origin(target_type)
|
|
check_args_exist(target_type)
|
|
if origin_type is list or origin_type is List:
|
|
if not isinstance(obj, list):
|
|
return False
|
|
arg_type = get_args(target_type)[0]
|
|
arg_origin = get_origin(arg_type)
|
|
for el in obj:
|
|
# check if nested container, ex: List[List[str]]
|
|
if arg_origin: # processes nested container, ex: List[List[str]]
|
|
if not container_checker(el, arg_type):
|
|
return False
|
|
elif not isinstance(el, arg_type):
|
|
return False
|
|
return True
|
|
elif origin_type is Dict or origin_type is dict:
|
|
if not isinstance(obj, dict):
|
|
return False
|
|
key_type = get_args(target_type)[0]
|
|
val_type = get_args(target_type)[1]
|
|
for key, val in obj.items():
|
|
# check if keys are of right type
|
|
if not isinstance(key, key_type):
|
|
return False
|
|
val_origin = get_origin(val_type)
|
|
if val_origin:
|
|
if not container_checker(val, val_type):
|
|
return False
|
|
elif not isinstance(val, val_type):
|
|
return False
|
|
return True
|
|
elif origin_type is Tuple or origin_type is tuple:
|
|
if not isinstance(obj, tuple):
|
|
return False
|
|
arg_types = get_args(target_type)
|
|
if len(obj) != len(arg_types):
|
|
return False
|
|
for el, el_type in zip(obj, arg_types):
|
|
el_origin = get_origin(el_type)
|
|
if el_origin:
|
|
if not container_checker(el, el_type):
|
|
return False
|
|
elif not isinstance(el, el_type):
|
|
return False
|
|
return True
|
|
elif origin_type is Union: # actually handles Optional Case
|
|
if obj is None: # check before recursion because None is always fine
|
|
return True
|
|
optional_type = get_args(target_type)[0]
|
|
optional_origin = get_origin(optional_type)
|
|
if optional_origin:
|
|
return container_checker(obj, optional_type)
|
|
elif isinstance(obj, optional_type):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _isinstance(obj, target_type) -> bool:
|
|
origin_type = get_origin(target_type)
|
|
if origin_type:
|
|
return container_checker(obj, target_type)
|
|
|
|
# Check to handle weird python type behaviors
|
|
# 1. python 3.6 returns None for origin of containers without
|
|
# contained type (intead of returning outer container type)
|
|
# 2. non-typed optional origin returns as none instead
|
|
# of as optional in 3.6-3.8
|
|
check_args_exist(target_type)
|
|
|
|
# handle non-containers
|
|
return isinstance(obj, target_type)
|
|
|
|
|
|
class _TensorExtractor(pickle.Pickler):
|
|
def __init__(self, *args, tensors: List[torch.Tensor], **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.tensors = tensors
|
|
|
|
def persistent_id(self, obj):
|
|
if isinstance(obj, torch.Tensor):
|
|
self.tensors.append(obj)
|
|
return ""
|
|
else:
|
|
return None
|
|
|
|
|
|
def _extract_tensors(obj):
|
|
r"""
|
|
This function is exclusively called from C++.
|
|
See ``torch/csrc/jit/python/python_ivalue.h``.
|
|
|
|
It extracts the tensors contained in the given object, through pickling.
|
|
"""
|
|
tensors: List[torch.Tensor] = []
|
|
extractor = _TensorExtractor(io.BytesIO(), protocol=-1, tensors=tensors)
|
|
extractor.dump(obj)
|
|
return tensors
|