from_torchvision

from_torchvision#

stable_pretraining.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:

    • 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. Special handling:

    • in_channels (int): Number of input channels. If provided for ResNet models, the first conv layer will be modified to accept this many channels. Default is 3.

Returns:

The neural network model.

Return type:

torch.nn.Module