lerobot代码解读(3)---.cache--2_calculate_validation_loss.py

原创已于 2025-03-14 15:14:25 修改 · 374 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能

于 2025-03-14 15:13:16 首次发布

#导入了数学计算库 math、深度学习库 torch，以及自定义的数据集类 LeRobotDataset 和 LeRobotDatasetMetadata，还有扩散策略模型类 DiffusionPolicy。

import math

import torch

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy

指定预训练策略模型的路径，使用 from_pretrained 方法加载模型，将模型设置为评估模式，并将模型移动到指定设备上。

def main():
    device = torch.device("cuda")#使用 CUDA 设备进行计算

    # Download the diffusion policy for pusht environment
    pretrained_policy_path = "lerobot/diffusion_pusht"
    # OR uncomment the following to evaluate a policy from the local outputs/train folder.
    # pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")

    policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
    policy.eval()
    policy.to(device)

#定义 delta_timestamps 字典，用于指定加载数据时不同观测和动作的时间戳。，"observation.image": [-0.1, 0.0] 表示加载当前帧前 0.1 秒和当前时刻的图像观测数据；"action" 列表指定了一系列时间戳的动作数据，用于计算损失。

   # Set up the dataset.
    delta_timestamps = {
        # Load the previous image and state at -0.1 seconds before current frame,
        # then load current image and state corresponding to 0.0 second.
        "observation.image": [-0.1, 0.0],
        "observation.state": [-0.1, 0.0],
        # Load the previous action (-0.1), the next action to be executed (0.0),
        # and 14 future actions with a 0.1 seconds spacing. All these actions will be
        # used to calculate the loss.
        "action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
    }

#LeRobotDatasetMetadata("lerobot/pusht")：加载数据集的元信息，获取数据集的总集数。

通过计算，将数据集划分为训练集（90%）和验证集（10%）。

LeRobotDataset：分别创建训练集和验证集的数据集对象，指定数据集的存储库 ID、使用的集数和时间戳。

打印训练集和验证集的集数和帧数信息，方便用户了解数据集的划分情况。

  # Load the last 10% of episodes of the dataset as a validation set.
    # - Load dataset metadata
    dataset_metadata = LeRobotDatasetMetadata("lerobot/pusht")
    # - Calculate train and val episodes
    total_episodes = dataset_metadata.total_episodes
    episodes = list(range(dataset_metadata.total_episodes))
    num_train_episodes = math.floor(total_episodes * 90 / 100)
    train_episodes = episodes[:num_train_episodes]
    val_episodes = episodes[num_train_episodes:]
    print(f"Number of episodes in full dataset: {total_episodes}")
    print(f"Number of episodes in training dataset (90% subset): {len(train_episodes)}")
    print(f"Number of episodes in validation dataset (10% subset): {len(val_episodes)}")
    # - Load train an val datasets
    train_dataset = LeRobotDataset(
        "lerobot/pusht", episodes=train_episodes, delta_timestamps=delta_timestamps
    )
    val_dataset = LeRobotDataset("lerobot/pusht", episodes=val_episodes, delta_timestamps=delta_timestamps)
    print(f"Number of frames in training dataset (90% subset): {len(train_dataset)}")
    print(f"Number of frames in validation dataset (10% subset): {len(val_dataset)}")

创建验证集的数据加载器

 # Create dataloader for evaluation.创建一个数据加载器，用于批量加载验证集的数据。
    val_dataloader = torch.utils.data.DataLoader(
        val_dataset,
        num_workers=4,使用 4 个工作线程来并行加载数据，提高数据加载效率
        batch_size=64,每个批次加载 64 个样本
        shuffle=False,不打乱数据顺序，保证每次加载的数据顺序一致。
        pin_memory=device != torch.device("cpu"),如果使用 GPU 计算，将数据预先加载到固定内存中，加快数据传输到 GPU 的速度
        drop_last=False,不丢弃最后一个不完整的批次。
    )

运行验证循环并计算平均损失

# Run validation loop.计算平均损失，运行验证循环。
    loss_cumsum = 0#用于累加每个批次的损失值。
    n_examples_evaluated = 0#用于记录评估的样本总数
    for batch in val_dataloader:
        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
        loss, _ = policy.forward(batch)# 调用模型的 forward 方法，计算当前批次的损失。

        loss_cumsum += loss.item()
        n_examples_evaluated += batch["index"].shape[0]# 累加当前批次的样本数。

    # Calculate the average loss over the validation set.
    average_loss = loss_cumsum / n_examples_evaluated#计算验证集的平均损失。

    print(f"Average loss on validation set: {average_loss:.4f}")


if __name__ == "__main__":
main()

这段代码的主要功能是加载预训练的扩散策略模型，并且将他们分为训练集和数据集，创建验证机的数据加载器。遍历验证集，计算平均损失。评估模型在验证集上的性能。