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]