A utility function to help freeze specific layers. Args: model: a source PyTorch model to freeze layer. freeze_vars: a regular expression to match the `model` variable names, so that their `requires_grad` will set to `False`. exclude_vars: a regular expr
(model: nn.Module, freeze_vars=None, exclude_vars=None)
| 1190 | |
| 1191 | |
| 1192 | def freeze_layers(model: nn.Module, freeze_vars=None, exclude_vars=None): |
| 1193 | """ |
| 1194 | A utility function to help freeze specific layers. |
| 1195 | |
| 1196 | Args: |
| 1197 | model: a source PyTorch model to freeze layer. |
| 1198 | freeze_vars: a regular expression to match the `model` variable names, |
| 1199 | so that their `requires_grad` will set to `False`. |
| 1200 | exclude_vars: a regular expression to match the `model` variable names, |
| 1201 | except for matched variable names, other `requires_grad` will set to `False`. |
| 1202 | |
| 1203 | Raises: |
| 1204 | ValueError: when freeze_vars and exclude_vars are both specified. |
| 1205 | |
| 1206 | """ |
| 1207 | if freeze_vars is not None and exclude_vars is not None: |
| 1208 | raise ValueError("Incompatible values: freeze_vars and exclude_vars are both specified.") |
| 1209 | src_dict = get_state_dict(model) |
| 1210 | |
| 1211 | frozen_keys = list() |
| 1212 | if freeze_vars is not None: |
| 1213 | to_freeze = {s_key for s_key in src_dict if freeze_vars and re.compile(freeze_vars).search(s_key)} |
| 1214 | for name, param in model.named_parameters(): |
| 1215 | if name in to_freeze: |
| 1216 | param.requires_grad = False |
| 1217 | frozen_keys.append(name) |
| 1218 | elif not param.requires_grad: |
| 1219 | param.requires_grad = True |
| 1220 | warnings.warn( |
| 1221 | f"The freeze_vars does not include {param}, but requires_grad is False, change it to True." |
| 1222 | ) |
| 1223 | if exclude_vars is not None: |
| 1224 | to_exclude = {s_key for s_key in src_dict if exclude_vars and re.compile(exclude_vars).search(s_key)} |
| 1225 | for name, param in model.named_parameters(): |
| 1226 | if name not in to_exclude: |
| 1227 | param.requires_grad = False |
| 1228 | frozen_keys.append(name) |
| 1229 | elif not param.requires_grad: |
| 1230 | param.requires_grad = True |
| 1231 | warnings.warn(f"The exclude_vars includes {param}, but requires_grad is False, change it to True.") |
| 1232 | |
| 1233 | logger.info(f"{len(frozen_keys)} of {len(src_dict)} variables frozen.") |
| 1234 | |
| 1235 | |
| 1236 | class CastTempType(nn.Module): |
searching dependent graphs…