mml.core.models.timm

class TimmGenericModel[source]

Bases: BaseModel

__init__(**kwargs)[source]
forward(x: Tensor) Dict[str, Tensor][source]
forward_features(x: Tensor) Tensor[source]
supports(task_type: TaskType) bool[source]

TimmModel supports classification and regression tasks.

class TimmHead[source]

Bases: BaseHead

__init__(task_type: TaskType, num_classes: int, num_features: int, drop_rate: float)[source]
forward(x: Tensor) Tensor[source]