create_model()

Create a classification model for cell image data.

Usage

Source

create_model(
    model,
    num_classes,
    num_channels=1,
    weights=None,
    freeze_weights=False,
    dropout=0.5
)

Parameters

model: Union[Model, str]

Model architecture enum value or string name.

num_classes: int

Number of output classes.

num_channels: int = 1

Number of input image channels.

weights: str | None = None

Pretrained weights identifier (e.g. 'DEFAULT').

freeze_weights: bool = False

Freeze pretrained weights; only new layers are trained.

dropout: float = 0.5

Dropout probability for the final classification layer.

Returns

nn.Module | None

Instantiated nn.Module, or None if the model is not recognised.