SwAVHook

class mmpretrain.engine.hooks.SwAVHook(batch_size, epoch_queue_starts=15, crops_for_assign=[0, 1], feat_dim=128, queue_length=0, interval=1, frozen_layers_cfg={})

Hook for SwAV.

This hook builds the queue in SwAV according to epoch_queue_starts. The queue will be saved in runner.work_dir or loaded at start epoch if the path folder has queues saved before.

Parameters:

• batch_size (int) – the batch size per GPU for computing.

• epoch_queue_starts (int, optional) – from this epoch, starts to use the queue. Defaults to 15.

• crops_for_assign (list[int], optional) – list of crops id used for computing assignments. Defaults to [0, 1].

• feat_dim (int, optional) – feature dimension of output vector. Defaults to 128.

• queue_length (int, optional) – length of the queue (0 for no queue). Defaults to 0.

• interval (int, optional) – the interval to save the queue. Defaults to 1.

• frozen_layers_cfg (dict, optional) – Dict to config frozen layers. The key-value pair is layer name and its frozen iters. If frozen, the layers don’t need gradient. Defaults to dict().

after_train_epoch(runner)

Save the queues locally.

before_run(runner)

Check whether the queues exist locally or not.

before_train_epoch(runner)

Check the queues’ state.

before_train_iter(runner, batch_idx, data_batch=None)

Freeze layers before specific iters according to the config.