Reference for ultralytics/models/yolo/pose/train.py

def get_dataset(self):
    """
    Retrieves the dataset and ensures it contains the required `kpt_shape` key.

    Returns:
        (dict): A dictionary containing the training/validation/test dataset and category names.

    Raises:
        KeyError: If the `kpt_shape` key is not present in the dataset.
    """
    data = super().get_dataset()
    if "kpt_shape" not in data:
        raise KeyError(f"No `kpt_shape` in the {self.args.data}. See https://docs.ultralytics.com/datasets/pose/")
    return data

def get_model(self, cfg=None, weights=None, verbose=True):
    """
    Get pose estimation model with specified configuration and weights.

    Args:
        cfg (str | Path | dict | None): Model configuration file path or dictionary.
        weights (str | Path | None): Path to the model weights file.
        verbose (bool): Whether to display model information.

    Returns:
        (PoseModel): Initialized pose estimation model.
    """
    model = PoseModel(
        cfg, nc=self.data["nc"], ch=self.data["channels"], data_kpt_shape=self.data["kpt_shape"], verbose=verbose
    )
    if weights:
        model.load(weights)

    return model

def get_validator(self):
    """Returns an instance of the PoseValidator class for validation."""
    self.loss_names = "box_loss", "pose_loss", "kobj_loss", "cls_loss", "dfl_loss"
    return yolo.pose.PoseValidator(
        self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
    )

def plot_metrics(self):
    """Plots training/val metrics."""
    plot_results(file=self.csv, pose=True, on_plot=self.on_plot)  # save results.png

def plot_training_samples(self, batch, ni):
    """
    Plot a batch of training samples with annotated class labels, bounding boxes, and keypoints.

    Args:
        batch (dict): Dictionary containing batch data with the following keys:
            - img (torch.Tensor): Batch of images
            - keypoints (torch.Tensor): Keypoints coordinates for pose estimation
            - cls (torch.Tensor): Class labels
            - bboxes (torch.Tensor): Bounding box coordinates
            - im_file (list): List of image file paths
            - batch_idx (torch.Tensor): Batch indices for each instance
        ni (int): Current training iteration number used for filename

    The function saves the plotted batch as an image in the trainer's save directory with the filename
    'train_batch{ni}.jpg', where ni is the iteration number.
    """
    images = batch["img"]
    kpts = batch["keypoints"]
    cls = batch["cls"].squeeze(-1)
    bboxes = batch["bboxes"]
    paths = batch["im_file"]
    batch_idx = batch["batch_idx"]
    plot_images(
        images,
        batch_idx,
        cls,
        bboxes,
        kpts=kpts,
        paths=paths,
        fname=self.save_dir / f"train_batch{ni}.jpg",
        on_plot=self.on_plot,
    )

def set_model_attributes(self):
    """Sets keypoints shape attribute of PoseModel."""
    super().set_model_attributes()
    self.model.kpt_shape = self.data["kpt_shape"]