Custom Models

class CustomModel

A named tuple that can be used to initialize a custom PyTorch model.

Args

• name: The name of the custom model. Default is 'default_model'.

• model: The PyTorch model object which is a subclass of torch.nn.Module and implements the forward method and output a tensor of shape (batch_size x features). Alternatively, a call method is also accepted.. Default is None.

• transform: A function that transforms a PIL.Image object into a PyTorch tensor that will be applied to each image before being fed to the model. Should correspond to the preprocessing logic of the supplied model. Default is None.

class MobilenetV3

__init__

def __init__()

Initialize a mobilenetv3 model, cuts it at the global average pooling layer and returns the output features.

forward

def forward(x)

class ViT

__init__

def __init__()

Initialize a ViT model, takes mean of the final encoder layer outputs and returns those as features for a given image.

forward

def forward(x)

class EfficientNet

__init__

def __init__()

Initializes an EfficientNet model, cuts it at the global average pooling layer and returns the output features.

forward

def forward(x)