1# mypy: allow-untyped-defs 2"""This module contains utility method for mobile model optimization and lint.""" 3 4import torch 5from enum import Enum 6from torch._C import _MobileOptimizerType as MobileOptimizerType 7from typing import Optional, Set, List, AnyStr 8 9class LintCode(Enum): 10 BUNDLED_INPUT = 1 11 REQUIRES_GRAD = 2 12 DROPOUT = 3 13 BATCHNORM = 4 14 15def 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 79def 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 120def _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