pytorchvideo.models.byol¶
-
class
pytorchvideo.models.byol.
BYOL
(backbone, projector=None, predictor=None, feature_dim=2048, predictor_inner=4096, mmt=0.99, norm=<class 'torch.nn.modules.batchnorm.SyncBatchNorm'>)[source]¶ Bootstrap Your Own Latent A New Approach to Self-Supervised Learning Details can be found in: https://arxiv.org/pdf/2006.07733.pdf
-
__init__
(backbone, projector=None, predictor=None, feature_dim=2048, predictor_inner=4096, mmt=0.99, norm=<class 'torch.nn.modules.batchnorm.SyncBatchNorm'>)[source]¶ - Parameters
backbone (nn.Module) – backbone for byol, input shape depends on the forward input size. Standard inputs include B x C, B x C x H x W, and B x C x T x H x W.
projector (nn.Module) – stand projector is a mlp with 2 to 3 hidden layers, with (synchronized) BatchNorm and ReLU activation.
predictor (nn.Module) – predictor MLP of BYOL of similar structure as the projector MLP.
feature_dim (int) – output feature dimension.
predictor_inner (int) – inner channel size for predictor.
mmt (float) – momentum update ratio for the momentum backbone.
norm (callable) – normalization to be used in projector, default is synchronized batchnorm.
- Return type
-
sim_loss
(q, k)[source]¶ Similarity loss for byol. :param q and k: inputs to calculate the similarity, expected to have
the same shape of N x C.
-
update_mmt
(mmt)[source]¶ Update the momentum. This function can be used to perform momentum annealing. :param mmt: update the momentum. :type mmt: float
- Parameters
mmt (float) –
-
get_mmt
()[source]¶ Get the momentum. This function can be used to perform momentum annealing.
- Return type
-
forward_backbone_mmt
(x)[source]¶ Forward momentum backbone. :param x: input to be forwarded. :type x: tensor
-
forward
(x1, x2)[source]¶ - Parameters
x1 (torch.tensor) – a batch of image with augmentation. The input tensor shape should able to be feed into the backbone.
x2 (torch.tensor) – the size batch of image with different augmentation. The input tensor shape should able to be feed into the backbone.
- Return type
-