U Ph%@sXddlmZddlmZmZddlZddlmZmZdgZ dddd d d d d ddZ dS)) annotations)CallableSequenceN) ensure_tupleensure_tuple_repgenerate_param_groupsTztorch.nn.ModulezSequence[Callable]z Sequence[str]zSequence[float]boolz list[dict])network layer_matches match_types lr_valuesinclude_othersreturnc st|}t|t|}t|t|}fdd}fdd}g}gt|||D]n\}} } | dkrp||} n&| dkr||} ntd| d|| | d d d | DqP|r|d tfd d i|S)a Utility function to generate parameter groups with different LR values for optimizer. The output parameter groups have the same order as `layer_match` functions. Args: network: source network to generate parameter groups from. layer_matches: a list of callable functions to select or filter out network layer groups, for "select" type, the input will be the `network`, for "filter" type, the input will be every item of `network.named_parameters()`. for "select", the parameters will be `select_func(network).parameters()`. for "filter", the parameters will be `(x[1] for x in filter(f, network.named_parameters()))` match_types: a list of tags to identify the matching type corresponding to the `layer_matches` functions, can be "select" or "filter". lr_values: a list of LR values corresponding to the `layer_matches` functions. include_others: whether to include the rest layers as the last group, default to True. It's mainly used to set different LR values for different network elements, for example: .. code-block:: python net = Unet(spatial_dims=3, in_channels=1, out_channels=3, channels=[2, 2, 2], strides=[1, 1, 1]) print(net) # print out network components to select expected items print(net.named_parameters()) # print out all the named parameters to filter out expected items params = generate_param_groups( network=net, layer_matches=[lambda x: x.model[0], lambda x: "2.0.conv" in x[0]], match_types=["select", "filter"], lr_values=[1e-2, 1e-3], ) # the groups will be a list of dictionaries: # [{'params': , 'lr': 0.01}, # {'params': , 'lr': 0.001}, # {'params': }] optimizer = torch.optim.Adam(params, 1e-4) csfdd}|S)Ncs SN) parametersfr rK/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/optimizers/utils.py_selectJsz;generate_param_groups.._get_select.._selectr)rrr rr _get_selectHsz*generate_param_groups.._get_selectcsfdd}|S)NcsddtDS)Ncss|]}|dVqdS)Nr.0xrrr SszNgenerate_param_groups.._get_filter.._filter..)filternamed_parametersrrrr_filterQsz;generate_param_groups.._get_filter.._filterr)rr rrr _get_filterOsz*generate_param_groups.._get_filterselectrzunsupported layer match type: .)paramslrcSsg|] }t|qSridrrrr bsz)generate_param_groups..r$cs t|kSrr&)p)_layersrrez'generate_param_groups..) rrlenziplower ValueErrorappendextendrr) r r r r r rr!r$functyr% layer_paramsr)r*r rrs$-       )T) __future__rcollections.abcrrtorch monai.utilsrr__all__rrrrr s