easycv.models.selfsup package

Submodules

easycv.models.selfsup.byol module

class easycv.models.selfsup.byol.BYOL(backbone, neck=None, head=None, pretrained=None, base_momentum=0.996, **kwargs)

Bases: easycv.models.base.BaseModel

BYOL unofficial implementation. Paper: https://arxiv.org/abs/2006.07733

__init__(backbone, neck=None, head=None, pretrained=None, base_momentum=0.996, **kwargs)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights()

forward_train(img, **kwargs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_test(img, **kwargs)

Abstract interface for model forward in testing

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward(img, mode='train', **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

easycv.models.selfsup.dino module

class easycv.models.selfsup.dino.MultiCropWrapper(backbone, head)

Bases: torch.nn.modules.module.Module

Perform forward pass separately on each resolution input. The inputs corresponding to a single resolution are clubbed and single forward is run on the same resolution inputs. Hence we do several forward passes = number of different resolutions used. We then concatenate all the output features and run the head forward on these concatenated features.

__init__(backbone, head)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.dino.DINOLoss(out_dim, ncrops, warmup_teacher_temp, teacher_temp, warmup_teacher_temp_epochs, nepochs, device, student_temp=0.1, center_momentum=0.9)

Bases: torch.nn.modules.module.Module

__init__(out_dim, ncrops, warmup_teacher_temp, teacher_temp, warmup_teacher_temp_epochs, nepochs, device, student_temp=0.1, center_momentum=0.9)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(student_output, teacher_output, epoch)

Cross-entropy between softmax outputs of the teacher and student networks.

update_center(teacher_output)

Update center used for teacher output.

training: bool

easycv.models.selfsup.dino.has_batchnorms(model)

easycv.models.selfsup.dino.get_params_groups(model)

class easycv.models.selfsup.dino.DINO(backbone, train_preprocess=[], neck=None, config=None, pretrained=None)

Bases: easycv.models.base.BaseModel

__init__(backbone, train_preprocess=[], neck=None, config=None, pretrained=None)

Init Moby

Parameters

• backbone – backbone config to build vision backbone

• train_preprocess – [gaussBlur, mixUp, solarize]

• neck – neck config to build Moby Neck

• config – DINO parameter config

get_params_groups()

init_weights(pretrained=None)

init_before_train()

momentum_update_key_encoder(m=0.999)

ema for dino

forward_train(inputs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_test(img, **kwargs)

Abstract interface for model forward in testing

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_feature(img, **kwargs)

Forward backbone

Returns

feature tensor

Return type

x (torch.Tensor)

training: bool

forward(img, gt_label=None, mode='train', extract_list=['neck'], **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

easycv.models.selfsup.mae module

class easycv.models.selfsup.mae.MAE(backbone, neck, mask_ratio=0.75, norm_pix_loss=True, **kwargs)

Bases: easycv.models.base.BaseModel

__init__(backbone, neck, mask_ratio=0.75, norm_pix_loss=True, **kwargs)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights()

patchify(imgs)

convert image to patch

Parameters

imgs – (N, 3, H, W)

Returns

(N, L, patch_size**2 *3)

Return type

x

forward_loss(imgs, pred, mask)

compute loss

Parameters

• imgs – (N, 3, H, W)

• pred – (N, L, p*p*3)

• mask – (N, L), 0 is keep, 1 is remove,

forward_train(img, **kwargs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_test(img, **kwargs)

Abstract interface for model forward in testing

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward(img, mode='train', **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

easycv.models.selfsup.mixco module

class easycv.models.selfsup.mixco.MIXCO(backbone, train_preprocess=[], neck=None, head=None, mixco_head=None, pretrained=None, queue_len=65536, feat_dim=128, momentum=0.999, **kwargs)

Bases: easycv.models.selfsup.moco.MOCO

MOCO.

A mixup version moco https://arxiv.org/pdf/2010.06300.pdf

__init__(backbone, train_preprocess=[], neck=None, head=None, mixco_head=None, pretrained=None, queue_len=65536, feat_dim=128, momentum=0.999, **kwargs)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward_train(img, **kwargs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

training: bool

easycv.models.selfsup.moby module

class easycv.models.selfsup.moby.MoBY(backbone, train_preprocess=[], neck=None, head=None, pretrained=None, queue_len=4096, contrast_temperature=0.2, momentum=0.99, online_drop_path_rate=0.2, target_drop_path_rate=0.0, **kwargs)

Bases: easycv.models.base.BaseModel

MoBY. Part of the code is borrowed from: https://github.com/SwinTransformer/Transformer-SSL/blob/main/models/moby.py.

__init__(backbone, train_preprocess=[], neck=None, head=None, pretrained=None, queue_len=4096, contrast_temperature=0.2, momentum=0.99, online_drop_path_rate=0.2, target_drop_path_rate=0.0, **kwargs)

Init Moby

Parameters

• backbone – backbone config to build vision backbone

• train_preprocess – [gaussBlur, mixUp, solarize]

• neck – neck config to build Moby Neck

• head – head config to build Moby Neck

• pretrained – pretrained weight for backbone

• queue_len – moby queue length

• contrast_temperature – contrastive_loss temperature

• momentum – ema target weights momentum

• online_drop_path_rate – for transformer based backbone, set online model drop_path_rate

• target_drop_path_rate – for transformer based backbone, set target model drop_path_rate

init_weights()

forward_backbone(img)

contrastive_loss(q, k, queue)

forward_train(img, **kwargs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_test(img, **kwargs)

Abstract interface for model forward in testing

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_feature(img, **kwargs)

Forward backbone

Returns

feature tensor

Return type

x (torch.Tensor)

forward(img, gt_label=None, mode='train', extract_list=['neck'], **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

easycv.models.selfsup.moby.concat_all_gather(tensor)

Performs all_gather operation on the provided tensors. * Warning *: torch.distributed.all_gather has no gradient.

easycv.models.selfsup.moco module

class easycv.models.selfsup.moco.MOCO(backbone, train_preprocess=[], neck=None, head=None, pretrained=None, queue_len=65536, feat_dim=128, momentum=0.999, **kwargs)

Bases: easycv.models.base.BaseModel

MOCO. Part of the code is borrowed from: https://github.com/facebookresearch/moco/blob/master/moco/builder.py.

__init__(backbone, train_preprocess=[], neck=None, head=None, pretrained=None, queue_len=65536, feat_dim=128, momentum=0.999, **kwargs)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights()

forward_backbone(img)

forward_train(img, **kwargs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_test(img, **kwargs)

Abstract interface for model forward in testing

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_feature(img, **kwargs)

Forward backbone

Returns

feature tensor

Return type

x (torch.Tensor)

forward(img, gt_label=None, mode='train', extract_list=['neck'], **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

easycv.models.selfsup.moco.concat_all_gather(tensor)

Performs all_gather operation on the provided tensors. * Warning *: torch.distributed.all_gather has no gradient.

easycv.models.selfsup.necks module

class easycv.models.selfsup.necks.DINONeck(in_dim, out_dim, use_bn=False, norm_last_layer=True, nlayers=3, hidden_dim=2048, bottleneck_dim=256)

Bases: torch.nn.modules.module.Module

__init__(in_dim, out_dim, use_bn=False, norm_last_layer=True, nlayers=3, hidden_dim=2048, bottleneck_dim=256)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.necks.MoBYMLP(in_channels=256, hid_channels=4096, out_channels=256, num_layers=2, with_avg_pool=True)

Bases: torch.nn.modules.module.Module

__init__(in_channels=256, hid_channels=4096, out_channels=256, num_layers=2, with_avg_pool=True)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

init_weights(init_linear='normal')

training: bool

class easycv.models.selfsup.necks.NonLinearNeckSwav(in_channels, hid_channels, out_channels, with_avg_pool=True, export=False)

Bases: torch.nn.modules.module.Module

The non-linear neck in byol: fc-syncbn-relu-fc

__init__(in_channels, hid_channels, out_channels, with_avg_pool=True, export=False)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights(init_linear='normal')

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.necks.NonLinearNeckV0(in_channels, hid_channels, out_channels, sync_bn=False, with_avg_pool=True)

Bases: torch.nn.modules.module.Module

The non-linear neck in ODC, fc-bn-relu-dropout-fc-relu

__init__(in_channels, hid_channels, out_channels, sync_bn=False, with_avg_pool=True)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights(init_linear='normal')

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.necks.NonLinearNeckV1(in_channels, hid_channels, out_channels, with_avg_pool=True)

Bases: torch.nn.modules.module.Module

The non-linear neck in MoCO v2: fc-relu-fc

__init__(in_channels, hid_channels, out_channels, with_avg_pool=True)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights(init_linear='normal')

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.necks.NonLinearNeckV2(in_channels, hid_channels, out_channels, with_avg_pool=True)

Bases: torch.nn.modules.module.Module

The non-linear neck in byol: fc-bn-relu-fc

__init__(in_channels, hid_channels, out_channels, with_avg_pool=True)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights(init_linear='normal')

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.necks.NonLinearNeckSimCLR(in_channels, hid_channels, out_channels, num_layers=2, with_avg_pool=True)

Bases: torch.nn.modules.module.Module

SimCLR non-linear neck.

Structure: fc(no_bias)-bn(has_bias)-[relu-fc(no_bias)-bn(no_bias)].

The substructures in [] can be repeated. For the SimCLR default setting, the repeat time is 1.

However, PyTorch does not support to specify (weight=True, bias=False).

It only support “affine” including the weight and bias. Hence, the second BatchNorm has bias in this implementation. This is different from the offical implementation of SimCLR.

Since SyncBatchNorm in pytorch<1.4.0 does not support 2D input, the input is

expanded to 4D with shape: (N,C,1,1). I am not sure if this workaround has no bugs. See the pull request here: https://github.com/pytorch/pytorch/pull/29626

Parameters

• in_channels – input channel number

• hid_channels – hidden channels

• out_channels – output channel number

• num_layers (int) – number of fc layers, it is 2 in the SimCLR default setting.

• with_avg_pool – output with average pooling

__init__(in_channels, hid_channels, out_channels, num_layers=2, with_avg_pool=True)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights(init_linear='normal')

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.necks.RelativeLocNeck(in_channels, out_channels, sync_bn=False, with_avg_pool=True)

Bases: torch.nn.modules.module.Module

Relative patch location neck: fc-bn-relu-dropout

__init__(in_channels, out_channels, sync_bn=False, with_avg_pool=True)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights(init_linear='normal')

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.necks.MAENeck(num_patches, embed_dim=768, patch_size=16, in_chans=3, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16, mlp_ratio=4.0, norm_layer=functools.partial(<class 'torch.nn.modules.normalization.LayerNorm'>, eps=1e-06))

Bases: torch.nn.modules.module.Module

MAE decoder

Parameters

• num_patches (int) – number of patches from encoder

• embed_dim (int) – encoder embedding dimension

• patch_size (int) – encoder patch size

• in_chans (int) – input image channels

• decoder_embed_dim (int) – decoder embedding dimension

• decoder_depth (int) – number of decoder layers

• decoder_num_heads (int) – Parallel attention heads

• mlp_ratio (float) – mlp ratio

• norm_layer – type of normalization layer

__init__(num_patches, embed_dim=768, patch_size=16, in_chans=3, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16, mlp_ratio=4.0, norm_layer=functools.partial(<class 'torch.nn.modules.normalization.LayerNorm'>, eps=1e-06))

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights()

training: bool

forward(x, ids_restore)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class easycv.models.selfsup.necks.FastConvMAENeck(num_patches, embed_dim=768, patch_size=16, in_channels=3, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16, mlp_ratio=4.0, norm_layer=functools.partial(<class 'torch.nn.modules.normalization.LayerNorm'>, eps=1e-06))

Bases: easycv.models.selfsup.necks.MAENeck

Fast ConvMAE decoder, refer to: https://github.com/Alpha-VL/FastConvMAE

Parameters

• num_patches (int) – number of patches from encoder

• embed_dim (int) – encoder embedding dimension

• patch_size (int) – encoder patch size

• in_channels (int) – input image channels

• decoder_embed_dim (int) – decoder embedding dimension

• decoder_depth (int) – number of decoder layers

• decoder_num_heads (int) – Parallel attention heads

• mlp_ratio (float) – mlp ratio

• norm_layer – type of normalization layer

training: bool

__init__(num_patches, embed_dim=768, patch_size=16, in_channels=3, decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16, mlp_ratio=4.0, norm_layer=functools.partial(<class 'torch.nn.modules.normalization.LayerNorm'>, eps=1e-06))

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights()

forward(x, ids_restore)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

easycv.models.selfsup.simclr module

class easycv.models.selfsup.simclr.SimCLR(backbone, train_preprocess=[], neck=None, head=None, pretrained=None)

Bases: easycv.models.base.BaseModel

__init__(backbone, train_preprocess=[], neck=None, head=None, pretrained=None)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights()

forward_backbone(img)

Forward backbone

Returns

backbone outputs

Return type

x (tuple)

forward_train(img, **kwargs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_test(img, **kwargs)

Abstract interface for model forward in testing

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward(img, mode='train', **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

easycv.models.selfsup.swav module

class easycv.models.selfsup.swav.SWAV(backbone, train_preprocess=[], neck=None, config=None, pretrained=None)

Bases: easycv.models.base.BaseModel

__init__(backbone, train_preprocess=[], neck=None, config=None, pretrained=None)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

init_weights()

forward_backbone(img)

forward_train_model(inputs)

forward_train(inputs)

Abstract interface for model forward in training

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_test(img, **kwargs)

Abstract interface for model forward in testing

Parameters

• img (Tensor) – image tensor

• kwargs (keyword arguments) – Specific to concrete implementation.

forward_feature(img, **kwargs)

Forward backbone

Returns

feature tensor

Return type

x (torch.Tensor)

forward(img, gt_label=None, mode='train', extract_list=['neck'], **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class easycv.models.selfsup.swav.MultiPrototypes(output_dim, nmb_prototypes)

Bases: torch.nn.modules.module.Module

__init__(output_dim, nmb_prototypes)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

easycv.models.selfsup.swav.distributed_sinkhorn(Q, nmb_iters)