FM / jit

/usr/local/lib64/python3.6/site-packages/torch/jit/

/usr/local/lib64/python3.6/site-packages/torch/jit

Name	Size	Mode	Actions
mobile/	-	0755	rm
__pycache__/	-	0755	rm
annotations.py	16126	0644	edit dl rm
frontend.py	40584	0644	edit dl rm
quantized.py	23816	0644	edit dl rm
supported_ops.py	10354	0644	edit dl rm
unsupported_tensor_ops.py	1830	0644	edit dl rm
_async.py	3787	0644	edit dl rm
_builtins.py	6205	0644	edit dl rm
_check.py	9280	0644	edit dl rm
_freeze.py	9047	0644	edit dl rm
_fuser.py	3819	0644	edit dl rm
_logging.py	256	0644	edit dl rm
_monkeytype_config.py	7129	0644	edit dl rm
_pickle.py	954	0644	edit dl rm
_recursive.py	40058	0644	edit dl rm
_script.py	60825	0644	edit dl rm
_serialization.py	6269	0644	edit dl rm
_state.py	3651	0644	edit dl rm
_trace.py	45873	0644	edit dl rm
__init__.py	6463	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/jit/_freeze.py (9047B)

"""Freezing

This is not intended to be imported directly; please use the exposed
functionalities in `torch.jit`.
"""

from typing import Optional, List

import torch
from torch.jit._script import RecursiveScriptModule, ScriptModule

def freeze(mod, preserved_attrs: Optional[List[str]] = None, optimize_numerics: bool = True):
    r"""
    Freezing a :class:`ScriptModule` will clone it and attempt to inline the cloned
    module's submodules, parameters, and attributes as constants in the TorchScript IR Graph.
    By default, `forward` will be preserved, as well as attributes & methods specified in
    `preserved_attrs`. Additionally, any attribute that is modified within a preserved
    method will be preserved.

Freezing currently only accepts ScriptModules that are in eval mode.

Freezing applies generic optimization that will speed up your model regardless of machine.
    To further optimize using server-specific settings, run `optimize_for_inference` after
    freezing.

Args:
        mod (:class:`ScriptModule`): a module to be frozen

preserved_attrs (Optional[List[str]]): a list of attributes to preserve in addition to the forward method.
        Attributes modified in preserved methods will also be preserved.

optimize_numerics (bool): If ``True``, a set of optimization passes will be run that does not strictly
        preserve numerics. Full details of optimization can be found at `torch.jit.run_frozen_optimizations`.

Returns:
        Frozen :class:`ScriptModule`.

Example (Freezing a simple module with a Parameter):

.. testcode::
        import torch
        class MyModule(torch.nn.Module):
            def __init__(self, N, M):
                super(MyModule, self).__init__()
                self.weight = torch.nn.Parameter(torch.rand(N, M))
                self.linear = torch.nn.Linear(N, M)

def forward(self, input):
                output = self.weight.mm(input)
                output = self.linear(output)
                return output

scripted_module = torch.jit.script(MyModule(2, 3).eval())
        frozen_module = torch.jit.freeze(scripted_module)
        # parameters have been removed and inlined into the Graph as constants
        assert len(list(frozen_module.named_parameters())) == 0
        # See the compiled graph as Python code
        print(frozen_module.code)

Example (Freezing a module with preserved attributes)

.. testcode::
        import torch
        class MyModule2(torch.nn.Module):
            def __init__(self):
                super(MyModule2, self).__init__()
                self.modified_tensor = torch.tensor(10.)
                self.version = 1

def forward(self, input):
                self.modified_tensor += 1
                return input + self.modified_tensor

scripted_module = torch.jit.script(MyModule2().eval())
        frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=["version"])
        # we've manually preserved `version`, so it still exists on the frozen module and can be modified
        assert frozen_module.version == 1
        frozen_module.version = 2
        # `modified_tensor` is detected as being mutated in the forward, so freezing preserves
        # it to retain model semantics
        assert frozen_module(torch.tensor(1)) == torch.tensor(12)
        # now that we've run it once, the next result will be incremented by one
        assert frozen_module(torch.tensor(1)) == torch.tensor(13)

Note:
        If you're not sure why an attribute is not being inlined as a constant, you can run
        `dump_alias_db` on frozen_module.forward.graph to see if freezing has detected the
        attribute is being modified.

Note:
        Because freezing makes weights constants and removes module hierarchy, `to` and other
        nn.Module methods to manipulate device or dtype no longer work. As a workaround,
        You can remap devices by specifying `map_location` in `torch.jit.load`, however
        device-specific logic may have been baked into the model.
    """
    if not isinstance(mod, ScriptModule):
        raise RuntimeError(
            "Freezing expects a ScriptModule as input. "
            "Please use torch.jit.script or torch.jit.trace to script your 'nn.Module'."
        )

if mod.training:
        raise RuntimeError(
            "Freezing is currently only implemented for modules in eval mode. "
            "Please call .eval() on your module before freezing."
        )

preserved_attrs = preserved_attrs if preserved_attrs is not None else []

out = RecursiveScriptModule(torch._C._freeze_module(mod._c, preserved_attrs))
    RecursiveScriptModule._finalize_scriptmodule(out)
    run_frozen_optimizations(out, optimize_numerics)

return out

def run_frozen_optimizations(mod, optimize_numerics: bool = True):
    r"""
    Runs a series of optimizations looking for patterns that occur in frozen graphs.
    The current set of optimizations is:
        - Dropout Removal
        - Conv -> Batchnorm folding
        - Conv -> Add/Sub folding
        - Conv -> Mul/Div folding

Args:
        mod (:class:`ScriptModule`): a frozen module to be optimized

optimize_numerics (bool): If ``True``, a set of optimization passes will be run that does not strictly
        preserve numerics. These optimizations preserve default rtol and atol of `torch.testing.assert_allclose`
        when applied on a single transformation, however in a module where many transformations are applied
        the rtol or atol may no longer fall within the default `assert_allclose` tolerance. Conv -> Batchnorm folding,
        Conv-Add/Sub, and Conv -> Mul/Div folding all may alter numerics.

Returns:
        None

Note:
        In rare occassions, this can result in slower execution.

Example (Freezing a module with Conv->Batchnorm)
    .. code-block:: python
        import torch
        in_channels, out_channels = 3, 32
        conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=True)
        bn = torch.nn.BatchNorm2d(out_channels, eps=.001)
        mod = torch.nn.Sequential(conv, bn)
        # set optimize to False here, by default freezing runs run_frozen_optimizations
        frozen_mod = torch.jit.freeze(torch.jit.script(mod.eval()), optimize=False)
        # inspect frozen mod
        assert "batch_norm" in str(frozen_mod.graph)
        torch.jit.run_frozen_optimizations(frozen_mod)
        assert "batch_norm" not in str(frozen_mod.graph)

"""
    # xxx: keep in sync with frozen_graph_optimization.cpp
    # intentionally duplicated to make to make it easier to create custom optimization sequence
    torch._C._jit_pass_remove_dropout(mod._c)
    if optimize_numerics:
        # run a couple times to capture Conv -> Mul -> Add etc
        for _ in range(2):
            torch._C._jit_pass_fold_frozen_conv_bn(mod.graph)
            torch._C._jit_pass_fold_frozen_conv_add_or_sub(mod.graph)
            torch._C._jit_pass_fold_frozen_conv_mul_or_div(mod.graph)

def optimize_for_inference(mod: ScriptModule) -> ScriptModule:
    """
    Performs a set of optimization passes to optimize a model for the
    purposes of inference. If the model is not already frozen, optimize_for_inference
    will invoke `torch.jit.freeze` automatically.

In addition to generic optimizations that should speed up your model regardless
    of environment, prepare for inference will also bake in build specific settings
    such as the presence of CUDNN or MKLDNN, and may in the future make transformations
    which speed things up on one machine but slow things down on another. Accordingly,
    serialization is not implemented following invoking `optimize_for_inference` and
    is not guaranteed.

This is still in prototype, and may have the potential to slow down your model.
    Primary use cases that have been targeted so far have been vision models on cpu
    and gpu to a lesser extent.

Example (optimizing a module with Conv->Batchnorm)::

import torch
        in_channels, out_channels = 3, 32
        conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=True)
        bn = torch.nn.BatchNorm2d(out_channels, eps=.001)
        mod = torch.nn.Sequential(conv, bn)
        frozen_mod = torch.jit.optimize_for_inference(torch.jit.script(mod.eval()))
        assert "batch_norm" not in str(frozen_mod.graph)
        # if built with MKLDNN, convolution will be run with MKLDNN weights
        assert "MKLDNN" in frozen_mod.graph
    """
    if not isinstance(mod, ScriptModule):
        raise RuntimeError(
            "optimize_for_inference expects a ScriptModule as input. "
            "Please use torch.jit.script or torch.jit.trace to script your 'nn.Module'.")

if hasattr(mod, "training"):
        mod = freeze(mod.eval())

torch._C._jit_pass_convert_frozen_ops_to_mkldnn(mod.graph)
    torch._C._jit_pass_fuse_frozen_conv_add_relu(mod.graph)
    return mod

cmd: