pytorch

1
# mypy: allow-untyped-defs
2
"""This module contains utility method for mobile model optimization and lint."""
3

4
import torch
5
from enum import Enum
6
from torch._C import _MobileOptimizerType as MobileOptimizerType
7
from typing import Optional, Set, List, AnyStr
8

9
class LintCode(Enum):
10
    BUNDLED_INPUT = 1
11
    REQUIRES_GRAD = 2
12
    DROPOUT = 3
13
    BATCHNORM = 4
14

15
def optimize_for_mobile(
16
        script_module: torch.jit.ScriptModule,
17
        optimization_blocklist: Optional[Set[MobileOptimizerType]] = None,
18
        preserved_methods: Optional[List[AnyStr]] = None,
19
        backend: str = 'CPU') -> torch.jit.RecursiveScriptModule:
20
    """
21
    Optimize a torch script module for mobile deployment.
22

23
    Args:
24
        script_module: An instance of torch script module with type of ScriptModule.
25
        optimization_blocklist: A set with type of MobileOptimizerType. When set is not passed,
26
            optimization method will run all the optimizer pass; otherwise, optimizer
27
            method will run the optimization pass that is not included inside optimization_blocklist.
28
        preserved_methods: A list of methods that needed to be preserved when freeze_module pass is invoked
29
        backend: Device type to use for running the result model ('CPU'(default), 'Vulkan' or 'Metal').
30
    Returns:
31
        A new optimized torch script module
32
    """
33
    if not isinstance(script_module, torch.jit.ScriptModule):
34
        raise TypeError(
35
            f'Got {type(script_module)}, but ScriptModule is expected.')
36

37
    if optimization_blocklist is None:
38
        optimization_blocklist = set()
39

40
    if preserved_methods is None:
41
        preserved_methods = []
42

43
    # Convert potential byte arrays into strings (if there is any) to pass type checking
44
    # Here we use a new name as assigning it back to preserved_methods will invoke
45
    # mypy errors (i.e. List[AnyStr] = List[str])
46
    preserved_methods_str: List[str] = [str(method) for method in preserved_methods]
47

48
    bundled_inputs_attributes = _get_bundled_inputs_preserved_attributes(script_module, preserved_methods_str)
49
    if all(hasattr(script_module, method) for method in bundled_inputs_attributes):
50
        preserved_methods_str = list(set(preserved_methods_str + bundled_inputs_attributes))
51

52
    non_exist_methods = []
53
    for method in preserved_methods_str:
54
        if not hasattr(script_module, method):
55
            non_exist_methods.append(method)
56
    if non_exist_methods:
57
        raise AttributeError(
58
            f"The following methods to preserve do not exist in script_module: {', '.join(non_exist_methods)}")
59

60
    backend = backend.lower()
61
    if backend == 'cpu':
62
        optimized_cpp_module = torch._C._jit_pass_optimize_for_mobile(
63
            script_module._c,
64
            optimization_blocklist,
65
            preserved_methods_str)
66
    elif backend == 'vulkan':
67
        optimized_cpp_module = torch._C._jit_pass_vulkan_optimize_for_mobile(
68
            script_module._c,
69
            optimization_blocklist,
70
            preserved_methods_str)
71
    elif backend == 'metal':
72
        optimized_cpp_module = torch._C._jit_pass_metal_optimize_for_mobile(script_module._c, preserved_methods_str)
73
    else:
74
        raise TypeError("Unknown backend, must be one of 'CPU', 'Vulkan' or 'Metal'")
75

76
    return torch.jit._recursive.wrap_cpp_module(optimized_cpp_module)
77

78

79
def generate_mobile_module_lints(script_module: torch.jit.ScriptModule):
80
    """
81
    Generate a list of lints for a given torch script module.
82

83
    Args:
84
        script_module: An instance of torch script module with type of ScriptModule.
85

86
    Returns:
87
        lint_map: A list of dictionary that contains modules lints
88
    """
89
    if not isinstance(script_module, torch.jit.ScriptModule):
90
        raise TypeError(
91
            f'Got {type(script_module)}, but ScriptModule is expected.')
92

93
    lint_list = []
94

95
    if not hasattr(script_module, "_generate_bundled_inputs_for_forward"):
96
        lint_list.append({"name": LintCode.BUNDLED_INPUT.name, "message": "No bundled input for forward, please add bundled inputs "
97
                          "before saving the module using torch.utils.bundled_inputs.augment_model_with_bundled_inputs."})
98

99
    for name, param in script_module.named_parameters():
100
        if param.requires_grad:
101
            lint_list.append({"name": LintCode.REQUIRES_GRAD.name, "message": f"Param {name} requires grad, "
102
                             "please set torch.no_grad() to reduce memory usage and improve computation speed during "
103
                              "inference phase."})
104

105
    op_names = torch.jit.export_opnames(script_module)
106
    for op_name in op_names:
107
        if "dropout" in op_name:
108
            lint_list.append({"name": LintCode.DROPOUT.name,
109
                              "message": f"Operator {op_name} exists, remember to call eval() before "
110
                              "saving the module.and call torch.utils.mobile_optimizer.optimize_for_mobile to drop dropout "
111
                              "operator."})
112
        if "batch_norm" in op_name:
113
            lint_list.append({"name": LintCode.BATCHNORM.name,
114
                              "message": f"Operator {op_name} exists, remember to call eval() before "
115
                              "saving the module and call torch.utils.mobile_optimizer.optimize_for_mobile to drop batch_norm "
116
                              "operator."})
117

118
    return lint_list
119

120
def _get_bundled_inputs_preserved_attributes(script_module: torch.jit.ScriptModule, preserved_methods: List[str]) -> List[str]:
121

122
    bundled_inputs_attributes = []
123
    # Has bundled inputs for forward
124
    if hasattr(script_module, 'get_all_bundled_inputs'):
125
        bundled_inputs_attributes.append('get_all_bundled_inputs')
126
        bundled_inputs_attributes.append('get_num_bundled_inputs')
127

128
    # Bundled inputs in module after the change that introduced bundled inputs for multiple functions
129
    if hasattr(script_module, 'get_bundled_inputs_functions_and_info'):
130
        bundled_inputs_attributes.append('get_bundled_inputs_functions_and_info')
131
        all_info = script_module.get_bundled_inputs_functions_and_info()
132
        for function_name in all_info:
133
            if function_name not in preserved_methods:
134
                bundled_inputs_attributes.append(function_name)
135
            bundled_inputs_attributes.append("get_all_bundled_inputs_for_" + function_name)
136
            bundled_inputs_attributes.append("_bundled_inputs_deflated_" + function_name)
137

138
    return bundled_inputs_attributes
139