relai.models package

Submodules

relai.models.cifar10_cnn module

class relai.models.cifar10_cnn.CIFAR10_CNN(num_classes: int = 10, transform=None)

Bases: RELAIModel

property backbone : Module

Model backbone; everything before the last fully connected layers.

property default_transform : Callable[[...], Any]

Default transform to be applied to the data before passing it to the model.

forward(x: Tensor)

Define 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.

property last_fc : Module

Last fully connected layers of the model.

relai.models.relai_model module

class relai.models.relai_model.RELAIModel(num_classes, transform=None, *args, **kwargs)

Bases: Module

abstract property backbone : Module

Model backbone; everything before the last fully connected layers.

abstract property default_transform : Callable[[...], Any]

Default transform to be applied to the data before passing it to the model.

finetune(mode: bool)

Set model to finetune mode or not.

abstract property last_fc : Module

Last fully connected layers of the model.

load_model_state(model_state)

property model_state

relai.models.resnet module

class relai.models.resnet.ResNet50(pretrained: bool = True, num_classes: int = None, transform: Callable[[...], Any] = None)

Bases: RELAIModel

ResNet50 model created from PyTorch’s torchvision.models.resnet50. See https://pytorch.org/vision/main/_modules/torchvision/models/resnet.html#resnet50.

• Parameters:
  • pretrained (bool , optional) – Whether to load pretrained weights. Defaults to True. When pretrained is set to True, the model is initialized with weights pretrained on ImageNet.
  • num_classes (int , optional) – Number of classes. Defaults to 1000. If num_classes is not 1000, the last layer is randomly initialized even if pretrained is True.

property backbone : Module

Model backbone; everything before the last fully connected layers.

property default_transform : Callable[[...], Any]

Default transform to be applied to the data before passing it to the model.

forward(x: Tensor)

Define 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.

property last_fc : Module

Last fully connected layers of the model.

relai.models.robust_resnet module

class relai.models.robust_resnet.RobustResNet50(mkey: str = None, path_to_models: str | Path = None, num_classes: int = None, transform: Callable[[...], Any] = None)

Bases: RELAIModel

Robust ResNet50 model pretrained on ImageNet. See https://github.com/MadryLab/robustness.

• Parameters:
  • num_classes (int , optional) – Number of classes. Defaults to 1000.
  • mkey (str) – Model name (e.g. “robust_resnet50_l2_eps0.25”).
  • path_to_models (Union *[*str , Path ]) – path to the robust models.
  • transform – transformation needed for this model.

property backbone : Module

Model backbone; everything before the last fully connected layers.

property default_transform : Callable[[...], Any]

Default transform to be applied to the data before passing it to the model.

forward(x: Tensor)

Define 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.

forward_features(x: Tensor)

property last_fc : Module

Last fully connected layers of the model.

relai.models.robust_resnet.load_pretrained_robust_model(mtype: str = 'robust_resnet50_l2_eps0.25', path_to_models: str | Path = None)

Returns a pretrained (on ImageNet) Robust ResNet model.

• Parameters:
  • mtype (str) – model type, in this format: robust_resnet50_l{n}_eps
  • path_to_models (Union *[*str , Path ]) – path to robust models.
• Return type: model (torch.nn.Module)

relai.models.robust_resnet.load_robust_model(num_classes: int, mkey: str = 'robust_resnet50_l2_eps0.25', path_to_models: str | Path = None)

Returns a pretrained robust model. This model must be later finetuned on a dataset. Note that training on a new dataset only effets the last linear layer.

• Parameters:
  • num_classes (int) – number of classes in the dataset on which this model is going to be evaluated on.
  • mtype (str) – model type, in this format: robust_resnet50_l{n}_eps
• Returns: a robust resnet model with forzen backbone and ready-to-learn last layer.
• Return type: model (torch.nn.Module)

relai.models.robust_resnet.make_and_restore_robust_model(*_, arch, dataset, resume_path=None, parallel=False, pytorch_pretrained=False, add_custom_forward=False, no_print=True)

relai.models.timm_model module

class relai.models.timm_model.TimmModel(pretrained: bool = True, num_classes: int = None, transform: Callable[[...], Any] = None, variant='resnet50', normalize_type='IMAGENET_DEFAULT')

Bases: RELAIModel

Timm models

• Parameters:
  • pretrained (bool , optional) – Whether to load pretrained weights. Defaults to True. When pretrained is set to True, the model is initialized with weights pretrained on ImageNet.
  • num_classes (int , optional) – Number of classes. Defaults to 1000. If num_classes is not 1000, the last layer is randomly initialized even if pretrained is True.

property default_transform : Callable[[...], Any]

Default transform to be applied to the data before passing it to the model.

forward(x: Tensor)

Define 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.

get_normalizer(normalize_type)

property last_fc : Module

Last fully connected layers of the model.

load_model_state(model_state)

property model_state

relai.models.timm_vit module

class relai.models.timm_vit.TimmVit(pretrained: bool = True, num_classes: int = None, transform: Callable[[...], Any] = None, variant='')

Bases: RELAIModel

Timm vit models

• Parameters:
  • pretrained (bool , optional) – Whether to load pretrained weights. Defaults to True. When pretrained is set to True, the model is initialized with weights pretrained on ImageNet.
  • num_classes (int , optional) – Number of classes. Defaults to 1000. If num_classes is not 1000, the last layer is randomly initialized even if pretrained is True.

property backbone : Module

Model backbone; everything before the last fully connected layers.

property default_transform : Callable[[...], Any]

Default transform to be applied to the data before passing it to the model.

forward(x: Tensor)

Define 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.

property last_fc : Module

Last fully connected layers of the model.

relai.models.utils module

relai.models.utils.get_detection_model(model_name, ckpt_path=None)

To load a detection model. Currently supporting detection models in mmdet format.

• Parameters:
  • model_name (str) – Name of model architecture to use. For mmdet model, this should start with a prefix of ‘mmdet:’, followed by either any model name supported by mmdet or path to a mmdet config file.
  • ckpt_path (Union *[*str , Path ] , optional) – Path to checkpoint to load. Defaults to None.
• Returns: For mmdet model,
• Raises:NotImplementedError – If the architecture specified in model_name is not implemented.

relai.models.utils.get_model(model_name: str, num_classes: int, transform=None, ckpt_path: str | Path = None)

Get model with specified architecture, number of classes and an optional checkpoint.

• Parameters:
  • model_name (str) – Name of model architecture to use.
  • num_classes (int) – Number of classes in dataset.
  • ckpt_path (Union *[*str , Path ] , optional) – Path to checkpoint to load. Defaults to None.
• Returns: Model with specified architecture and number of classes.
• Return type: torch.nn.Module
• Raises:NotImplementedError – If the architecture specified in model_name is not implemented.

relai.models.utils.modify_user_checkpoint_keys(model_name: str, ckpt: dict)

Modify user uploaded checkpoint keys to match to those expected by RELAIModel.

relai.models.utils.validate_checkpoint(model_name: str, ckpt_path: str | Path, dataset: RELAIClassificationDataset, is_relai_checkpoint: bool = False)

Validates a checkpoint by determining if it can be loaded into a model with the specified architecture and number of classes in the dataset.

• Parameters:
  • model_name (str) – Name of model architecture to use.
  • ckpt_path (Union *[*str , Path ]) – Path to checkpoint to load.
  • dataset (RELAIClassificationDataset) – Dataset to use for validation.
  • is_relai_checkpoint (bool , optional) – Whether the checkpoint is a RELAI checkpoint as opposed to an user uploaded one. Defaults to False.
• Returns: Whether the checkpoint has been modified. dict: A valid RELAI checkpoint.
• Return type: bool
• Raises:InvalidCheckpointRELAIError – If the checkpoint is not valid.

Module contents