1
"""This module contains utility method for mobile model optimization and lint."""
5
from torch._C import _MobileOptimizerType as MobileOptimizerType
6
from typing import Optional, Set, List, AnyStr
14
def optimize_for_mobile(
15
script_module: torch.jit.ScriptModule,
16
optimization_blocklist: Optional[Set[MobileOptimizerType]] = None,
17
preserved_methods: Optional[List[AnyStr]] = None,
18
backend: str = 'CPU') -> torch.jit.RecursiveScriptModule:
20
Optimize a torch script module for mobile deployment.
23
script_module: An instance of torch script module with type of ScriptModule.
24
optimization_blocklist: A set with type of MobileOptimizerType. When set is not passed,
25
optimization method will run all the optimizer pass; otherwise, optimizer
26
method will run the optimization pass that is not included inside optimization_blocklist.
27
preserved_methods: A list of methods that needed to be preserved when freeze_module pass is invoked
28
backend: Device type to use for running the result model ('CPU'(default), 'Vulkan' or 'Metal').
30
A new optimized torch script module
32
if not isinstance(script_module, torch.jit.ScriptModule):
34
f'Got {type(script_module)}, but ScriptModule is expected.')
36
if optimization_blocklist is None:
37
optimization_blocklist = set()
39
if preserved_methods is None:
40
preserved_methods = []
45
preserved_methods_str: List[str] = [str(method) for method in preserved_methods]
47
bundled_inputs_attributes = _get_bundled_inputs_preserved_attributes(script_module, preserved_methods_str)
48
if all(hasattr(script_module, method) for method in bundled_inputs_attributes):
49
preserved_methods_str = list(set(preserved_methods_str + bundled_inputs_attributes))
51
non_exist_methods = []
52
for method in preserved_methods_str:
53
if not hasattr(script_module, method):
54
non_exist_methods.append(method)
57
f"The following methods to preserve do not exist in script_module: {', '.join(non_exist_methods)}")
59
backend = backend.lower()
61
optimized_cpp_module = torch._C._jit_pass_optimize_for_mobile(
63
optimization_blocklist,
64
preserved_methods_str)
65
elif backend == 'vulkan':
66
optimized_cpp_module = torch._C._jit_pass_vulkan_optimize_for_mobile(
68
optimization_blocklist,
69
preserved_methods_str)
70
elif backend == 'metal':
71
optimized_cpp_module = torch._C._jit_pass_metal_optimize_for_mobile(script_module._c, preserved_methods_str)
73
raise TypeError("Unknown backend, must be one of 'CPU', 'Vulkan' or 'Metal'")
75
return torch.jit._recursive.wrap_cpp_module(optimized_cpp_module)
78
def generate_mobile_module_lints(script_module: torch.jit.ScriptModule):
80
Generate a list of lints for a given torch script module.
83
script_module: An instance of torch script module with type of ScriptModule.
86
lint_map: A list of dictionary that contains modules lints
88
if not isinstance(script_module, torch.jit.ScriptModule):
90
f'Got {type(script_module)}, but ScriptModule is expected.')
94
if not hasattr(script_module, "_generate_bundled_inputs_for_forward"):
95
lint_list.append({"name": LintCode.BUNDLED_INPUT.name, "message": "No bundled input for forward, please add bundled inputs "
96
"before saving the module using torch.utils.bundled_inputs.augment_model_with_bundled_inputs."})
98
for name, param in script_module.named_parameters():
99
if param.requires_grad:
100
lint_list.append({"name": LintCode.REQUIRES_GRAD.name, "message": f"Param {name} requires grad, "
101
"please set torch.no_grad() to reduce memory usage and improve computation speed during "
104
op_names = torch.jit.export_opnames(script_module)
105
for op_name in op_names:
106
if "dropout" in op_name:
107
lint_list.append({"name": LintCode.DROPOUT.name, "message": "Operator {} exists, remember to call eval() before "
108
"saving the module.and call torch.utils.mobile_optimizer.optimize_for_mobile to drop dropout "
109
"operator.".format(op_name)})
110
if "batch_norm" in op_name:
111
lint_list.append({"name": LintCode.BATCHNORM.name, "message": "Operator {} exists, remember to call eval() before "
112
"saving the module and call torch.utils.mobile_optimizer.optimize_for_mobile to drop batch_norm "
113
"operator.".format(op_name)})
117
def _get_bundled_inputs_preserved_attributes(script_module: torch.jit.ScriptModule, preserved_methods: List[str]) -> List[str]:
119
bundled_inputs_attributes = []
121
if hasattr(script_module, 'get_all_bundled_inputs'):
122
bundled_inputs_attributes.append('get_all_bundled_inputs')
123
bundled_inputs_attributes.append('get_num_bundled_inputs')
126
if hasattr(script_module, 'get_bundled_inputs_functions_and_info'):
127
bundled_inputs_attributes.append('get_bundled_inputs_functions_and_info')
128
all_info = script_module.get_bundled_inputs_functions_and_info()
129
for function_name in all_info:
130
if function_name not in preserved_methods:
131
bundled_inputs_attributes.append(function_name)
132
bundled_inputs_attributes.append("get_all_bundled_inputs_for_" + function_name)
133
bundled_inputs_attributes.append("_bundled_inputs_deflated_" + function_name)
135
return bundled_inputs_attributes