Add custom encoder argument to BYOL

[Luab]

🚀 Feature

Right now the BYOL module uses the vanilla resnet50 as encoder, although underlying SiameseArm module supports custom encoder. Passing encoder as argument to BYOL module would allow custom encoders to be used.

class BYOL(pl.LightningModule):
    def __init__(self,
                 num_classes,
                 base_encoder: Union[str, torch.nn.Module] = 'resnet50',
                 learning_rate: float = 0.2,
                 weight_decay: float = 1.5e-6,
                 input_height: int = 32,
                 batch_size: int = 32,
                 num_workers: int = 0,
                 warmup_epochs: int = 10,
                 max_epochs: int = 1000,
                 **kwargs):
        super().__init__()
        self.save_hyperparameters()

        self.online_network = SiameseArm(encoder=base_encoder)

Also small fix in SiameseArm module is needed

class SiameseArm(nn.Module):
    def __init__(self, encoder='resnet50'):
        super().__init__()

        if isinstance(encoder, str):
            encoder = torchvision_ssl_encoder(encoder)

Motivation

Adding encoder as an argument in the init method would allow seamless ablation studies by trying out different architectures.

[github-actions]

Hi! thanks for your contribution!, great first issue!