BaseSupervisedTracker

class lightning_pose.models.base.BaseSupervisedTracker(backbone: Literal['resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnet50_contrastive', 'resnet50_animal_apose', 'resnet50_animal_ap10k', 'resnet50_human_jhmdb', 'resnet50_human_res_rle', 'resnet50_human_top_res', 'resnet50_human_hand', 'efficientnet_b0', 'efficientnet_b1', 'efficientnet_b2', 'vit_b_sam'] = 'resnet50', pretrained: bool = True, lr_scheduler: str = 'multisteplr', lr_scheduler_params: DictConfig | dict | None = None, do_context: bool = False, image_size: int = 256, model_type: Literal['heatmap', 'regression'] = 'heatmap', **kwargs: Any)[source]

Bases: BaseFeatureExtractor

Base class for supervised trackers.

Methods Summary

`evaluate_labeled`(batch_dict[, stage])	Compute and log the losses on a batch of labeled data.
`get_loss_inputs_labeled`(batch_dict)	Return predicted coordinates for a batch of data.
`test_step`(batch_dict, batch_idx)	Base test step, a wrapper around the evaluate_labeled method.
`training_step`(batch_dict, batch_idx)	Base training step, a wrapper around the evaluate_labeled method.
`validation_step`(batch_dict, batch_idx)	Base validation step, a wrapper around the evaluate_labeled method.

Methods Documentation

evaluate_labeled(batch_dict: BaseLabeledBatchDict | HeatmapLabeledBatchDict | MultiviewLabeledBatchDict | MultiviewHeatmapLabeledBatchDict, stage: Literal['train', 'val', 'test'] | None = None) → Tensor[Tensor][source]: Compute and log the losses on a batch of labeled data.

get_loss_inputs_labeled(batch_dict: BaseLabeledBatchDict | HeatmapLabeledBatchDict | MultiviewLabeledBatchDict | MultiviewHeatmapLabeledBatchDict) → dict[source]: Return predicted coordinates for a batch of data.

test_step(batch_dict: BaseLabeledBatchDict | HeatmapLabeledBatchDict | MultiviewLabeledBatchDict | MultiviewHeatmapLabeledBatchDict, batch_idx: int) → None[source]: Base test step, a wrapper around the evaluate_labeled method.

training_step(batch_dict: BaseLabeledBatchDict | HeatmapLabeledBatchDict | MultiviewLabeledBatchDict | MultiviewHeatmapLabeledBatchDict, batch_idx: int) → Dict[str, Tensor[Tensor]][source]: Base training step, a wrapper around the evaluate_labeled method.

validation_step(batch_dict: BaseLabeledBatchDict | HeatmapLabeledBatchDict | MultiviewLabeledBatchDict | MultiviewHeatmapLabeledBatchDict, batch_idx: int) → None[source]: Base validation step, a wrapper around the evaluate_labeled method.