Models

Creating models

create_model(model_name, pretrained=False, model_path='', *, in_channels=None, nb_classes=None, **kwargs)[source]

Creates a model.

Parameters:

  • model_name (str) – Name of model to instantiate

  • pretrained (bool) –

    If True, load pretrained weights as specified by the url field in config. We will check the cache first and download weights only if they cannot be found in the cache.

    If url is [timm], the weights will be downloaded from timm and converted to TensorFlow. Requires timm and torch to be installed. If url starts with [pytorch], the weights are in PyTorch format and torch needs to be installed to convert them.

  • model_path (str) – Path of model weights to load after model is initialized. This takes over pretrained.

  • in_channels (int | None) – Number of input channels for model. If None, use default provided by model.

  • nb_classes (int | None) – Number of classes for classifier. If set to 0, no classifier is used and last layer is pooling layer. If None, use default provided by model.

  • **kwargs – Other kwargs are model specific.

Returns:

The created model.

Return type:

Model

create_preprocessing(model_name, *, in_channels=None, dtype=None)[source]

Creates a function to preprocess images for a particular model.

The input to the preprocessing function is assumed to be values in range [0, 255].

Parameters:

  • model_name (str) – Model for which to create preprocessing function.

  • in_channels (float | None) – Number of input channels to model

  • dtype (str | None) – Output dtype.

Returns:

Callable that operates on single images as well as batches.

Return type:

Callable