from_torchvision

from_torchvision#

stable_ssl.backbone.from_torchvision(model_name, low_resolution=False, **kwargs)[source]#

Load a backbone model.

If num_classes is provided, the last layer is replaced by a linear layer of output size num_classes. Otherwise, the last layer is replaced by an identity layer.

Parameters:

  • model_name (str) – Name of the backbone model. Supported models are: - Any model from torchvision.models - “Resnet9” - “ConvMixer”

  • low_resolution (bool, optional) – Whether to adapt the resolution of the model (for CIFAR typically). By default False.

  • **kwargs – Additional keyword arguments for the model.

Returns:

The neural network model.

Return type:

torch.nn.Module

Contents