fastreid.engine.hooks.FreezeLayer Class Reference

Inheritance diagram for fastreid.engine.hooks.FreezeLayer:

Public Member Functions
	__init__ (self, model, optimizer, freeze_layers, freeze_iters)
	before_step (self)
	freeze_specific_layer (self)
	open_all_layer (self)
Public Member Functions inherited from fastreid.engine.train_loop.HookBase
	before_train (self)
	after_train (self)
	after_step (self)

Public Attributes
	model
	optimizer
	freeze_layers
	freeze_iters
	param_freeze
	is_frozen

Protected Attributes
	_logger

Detailed Description

Definition at line 417 of file hooks.py.

Constructor & Destructor Documentation

__init__()

fastreid.engine.hooks.FreezeLayer.__init__	(		self,
			model,
			optimizer,
			freeze_layers,
			freeze_iters )

Definition at line 418 of file hooks.py.

    def __init__(self, model, optimizer, freeze_layers, freeze_iters):
        self._logger = logging.getLogger(__name__)
 
        if isinstance(model, DistributedDataParallel):
            model = model.module
        self.model = model
        self.optimizer = optimizer
 
        self.freeze_layers = freeze_layers
        self.freeze_iters = freeze_iters
 
        # Previous parameters freeze status
        param_freeze = {}
        for param_group in self.optimizer.param_groups:
            param_name = param_group['name']
            param_freeze[param_name] = param_group['freeze']
        self.param_freeze = param_freeze
 
        self.is_frozen = False
 

Member Function Documentation

before_step()

fastreid.engine.hooks.FreezeLayer.before_step

(

self

)

Called before each iteration.

Reimplemented from fastreid.engine.train_loop.HookBase.

Definition at line 438 of file hooks.py.

    def before_step(self):
        # Freeze specific layers
        if self.trainer.iter <= self.freeze_iters and not self.is_frozen:
            self.freeze_specific_layer()
 
        # Recover original layers status
        if self.trainer.iter > self.freeze_iters and self.is_frozen:
            self.open_all_layer()
 

freeze_specific_layer()

fastreid.engine.hooks.FreezeLayer.freeze_specific_layer

(

self

)

Definition at line 447 of file hooks.py.

    def freeze_specific_layer(self):
        for layer in self.freeze_layers:
            if not hasattr(self.model, layer):
                self._logger.info(f'{layer} is not an attribute of the model, will skip this layer')
 
        for param_group in self.optimizer.param_groups:
            param_name = param_group['name']
            if param_name.split('.')[0] in self.freeze_layers:
                param_group['freeze'] = True
 
        # Change BN in freeze layers to eval mode
        for name, module in self.model.named_children():
            if name in self.freeze_layers: module.eval()
 
        self.is_frozen = True
 

open_all_layer()

fastreid.engine.hooks.FreezeLayer.open_all_layer

(

self

)

Definition at line 463 of file hooks.py.

    def open_all_layer(self):
        self.model.train()
        for param_group in self.optimizer.param_groups:
            param_name = param_group['name']
            param_group['freeze'] = self.param_freeze[param_name]
 
        self.is_frozen = False
 
 

Member Data Documentation

_logger

fastreid.engine.hooks.FreezeLayer._logger

protected

Definition at line 419 of file hooks.py.

freeze_iters

fastreid.engine.hooks.FreezeLayer.freeze_iters

Definition at line 427 of file hooks.py.

freeze_layers

fastreid.engine.hooks.FreezeLayer.freeze_layers

Definition at line 426 of file hooks.py.

is_frozen

fastreid.engine.hooks.FreezeLayer.is_frozen

Definition at line 436 of file hooks.py.

model

fastreid.engine.hooks.FreezeLayer.model

Definition at line 423 of file hooks.py.

optimizer

fastreid.engine.hooks.FreezeLayer.optimizer

Definition at line 424 of file hooks.py.

param_freeze

fastreid.engine.hooks.FreezeLayer.param_freeze

Definition at line 434 of file hooks.py.

---

The documentation for this class was generated from the following file:

• smreid/fastreid/engine/hooks.py