a 7Sicx@sdZddlZddlmZddlmZddlmZmZm Z m Z GdddeZ dej j eeeee e eej jd d d Zej j d d dZej j e ee edddZdS)zM This module contains utility method for mobile model optimization and lint. N)Enum)MobileOptimizerType)OptionalSetListAnyStrc@seZdZdZdZdZdZdS)LintCodeN)__name__ __module__ __qualname__ BUNDLED_INPUT REQUIRES_GRADDROPOUT BATCHNORMrrX/var/www/html/django/DPS/env/lib/python3.9/site-packages/torch/utils/mobile_optimizer.pyr srCPU) script_moduleoptimization_blocklistpreserved_methodsbackendreturnc s"ttjjs tdt|dur.t}|dur:g}dd|D}t|}t fdd|Drxt t||}g}|D]}t |s| |q|rt dd||}|dkrtjj||}n>|d krtjj|}n$|d kr tjj|}ntd tjj|S) a Args: script_module: An instance of torch script module with type of ScriptModule. optimization_blocklist: A set with type of MobileOptimizerType. When set is not passed, optimization method will run all the optimizer pass; otherwise, optimizer method will run the optimization pass that is not included inside optimization_blocklist. preserved_methods: A list of methods that needed to be preserved when freeze_module pass is invoked backend: Device type to use for running the result model ('CPU'(default), 'Vulkan' or 'Metal'). Returns: A new optimized torch script module %Got {}, but ScriptModule is expected.NcSsg|] }t|qSr)str.0methodrrr -z'optimize_for_mobile..csg|]}t|qSr)hasattrrrrrr!0r"zCThe following methods to preserve do not exist in script_module: {}z, cpuZvulkanZmetalz:Unknown backend, must be one of 'CPU', 'Vulkan' or 'Metal') isinstancetorchjit ScriptModule TypeErrorformattypeset(_get_bundled_inputs_preserved_attributesalllistr#appendAttributeErrorjoinlower_C_jit_pass_optimize_for_mobile_c$_jit_pass_vulkan_optimize_for_mobile#_jit_pass_metal_optimize_for_mobile _recursivewrap_cpp_module) rrrrZpreserved_methods_strbundled_inputs_attributesZnon_exist_methodsr Zoptimized_cpp_modulerr$roptimize_for_mobilesF     r=r$cCst|tjjs tdt|g}t|dsB|t j j dd| D](\}}|j rJ|t jj d|dqJtj|}|D]H}d|vr|t jj d|dd|vr|t jj d |dq|S) z Args: script_module: An instance of torch script module with type of ScriptModule Returns: lint_map: A list of dictionary that contains modules lints rZ$_generate_bundled_inputs_for_forwardzNo bundled input for forward, please add bundled inputs before saving the module using torch.utils.bundled_inputs.augment_model_with_bundled_inputs.)namemessagezParam {} requires grad, please set torch.no_grad() to reduce memory usage and improve computation speed during inference phase.dropoutzOperator {} exists, remember to call eval() before saving the module.and call torch.utils.mobile_optimizer.optimize_for_mobile to drop dropout operator. batch_normzOperator {} exists, remember to call eval() before saving the module and call torch.utils.mobile_optimizer.optimize_for_mobile to drop batch_norm operator.)r&r'r(r)r*r+r,r#r1rrr>named_parameters requires_gradrexport_opnamesrr)rZ lint_listr>paramop_namesop_namerrrgenerate_mobile_module_lintsLs.      rH)rrrcCszg}t|dr"|d|dt|drv|d|}|D]2}||vrX|||d||d|qB|S)NZget_all_bundled_inputsZget_num_bundled_inputs%get_bundled_inputs_functions_and_infoZget_all_bundled_inputs_for_Z_bundled_inputs_deflated_)r#r1rI)rrr<Zall_info function_namerrrr.qs      r.)NNr)__doc__r'enumrtorch._Crtypingrrrrrr(r)rRecursiveScriptModuler=rHr.rrrrs      <%