gy77/add freeze hook

[gy-7]

Motivation

Motivation:

1. Freeze some parameters of the model when training the model.

Goal:

1. Specify the epoch to freeze the specified network layer.
2. Available for all downstream repositories.

Modification

Add FreezeHook and FreezeHook unit tests.

Use cases

1. Network layers matching freeze_layers are freeze before freeze_iter/freeze_epoch starts.
2. Network layers matching unfreeze_layers are freeze before unfreeze_iter/unfreeze_epoch starts.
3. freeze_layers/unfreeze_layers matches network layers via regular expression
4. The index of iter/epoch starts at 0, with epoch=0 for the first epoch.
5. unfreeze_iter, unfreeze_epoch and unfreeze_layers are optional. If freeze_epoch/freeze_iter is not None, unfreeze_layers must not be None.
6. Only one of freeze_iter and freeze_epoch can be set, as well as unfreeze_iter and unfreeze_epoch.

ImageClassifier(
    (backbone):ResNet(
        ...
        (layer1):Sequential(...)
        (layer2):Sequential(...)
        (layer3):Sequential(...)
        (layer4):Sequential(...)
    )
    (neck):GlobalAveragePooling2d(...)
    (head):Linear(...)
)

1. Freeze the parameters of backbone before the start of 1st training epoch.

custom_hooks = [
...
dict(
    type="FreezeHook",
    freeze_layers="backbone.*",
    freeze_epoch=0)
]

2. Freeze the layer1 and layer2 parameters in the backbone before the start of 10th training epoch.

custom_hooks = [
...
dict(
    type="FreezeHook",
    freeze_layers="backbone.layer1.*|backbone.layer2.*",
    freeze_epoch=10)
]

3. Freeze the parameters of backbone before the start of 1st training epoch. Unfreeze the parameters of the the backbone before the start of 10th training epoch.

custom_hooks = [
 ...
 dict(
     type="FreezeHook",
     freeze_layers="backbone.*",
     freeze_epoch=0,
     unfreeze_layers="backbone.*",
     unfreeze_epoch=9)
]

4. The verbose parameter is used to determine whether to print the requires_grad variable for each model layer.

custom_hooks = [
 ...
 dict(
     type="FreezeHook",
     freeze_layers="backbone.*",
     freeze_epoch=1,
     verbose=True)
]

mmengine - INFO - backbone.conv1.weight requires_grad: True
mmengine - INFO - backbone.bn1.weight requires_grad: True
...
mmengine - INFO - head.light_head.weight requires_grad: True
mmengine - INFO - head.light_head.bias requires_grad: True

[HAOCHENYE]

Suggested change

	log_grad (bool): Whether to log the requires_grad of each layer.
	verbose (bool): Whether to log the requires_grad of each layer.

[HAOCHENYE]

Suggest making freeze_layers the first argument, and it should be a regex expression

[HAOCHENYE]

Suggested change

	def modify_layers_grad(self, model, layers, requires_grad):
	def _modify_layers_grad(self, model, layers, requires_grad):

[HAOCHENYE]

Please update the type hint

[HAOCHENYE]

Suggested change

	def log_model_grad(self, model, log_grad=False):
	def _log_model_grad(self, model, log_grad=False):

[HAOCHENYE]

Since it has been a regex expression, it is not necessary to make it a tuple of str ('exp1|exp2|exp3' is enough)

[HAOCHENYE]

Suggested change

	if not isinstance(freeze_iter, (int, type(None))):
	if not isinstance(freeze_iter) and freeze_iter is not None:

[HAOCHENYE]

Suggested change

	if freeze_iter and freeze_iter < 0:
	if freeze_iter is not None and freeze_iter < 0:

[HAOCHENYE]

Suggested change

	if freeze_epoch and freeze_epoch < 0:
	if freeze_epoch is not None and freeze_epoch < 0:

[HAOCHENYE]

Merge this check into:

if (freeze_iter is None)  ^ (freeze_epoch is None):
   raise ValueError(...)

if freeze_iter is not None and freeze_iter < 0:
    raise ValueError(...)

if freeze_epoch is not None and freeze_epoch < 0:
    raise ValueError(...)

[HAOCHENYE]

Suggested change

	model (BaseModel): a BaseModel of mmengine.
	model (BaseModel): A BaseModel of mmengine.

[HAOCHENYE]

Suggested change

	"""Modify the `requires_grad` of the specified layers.
	"""Modify the ``requires_grad`` of the specified layers.

[HAOCHENYE]

Suggested change

	def _main(self,
	def _freeze(self,

[HAOCHENYE]

Suggested change

	if self.freeze_iter is not None:
	self._main(runner, runner.iter, self.freeze_iter,
	self.unfreeze_iter)
	if self.freeze_iter is not None and runner.iter in (self.freeze_iter, self.unfreeze_iter):
	self._freeze(runner.model)