Data loading is slow due to grabbing frames from huggingface dataset

[alexander-soare]

I ran an experiment where I benchmarked times for running a dataloader on xarm_lift_medium_replay with batch size 256 and 0 workers.

8 batches takes ~ 15s. Here's the breakdown:

• dataset.__getitem__: 14.707566491064426
• load_previous_and_future_frames: 13.887129129978348
• hf_dataset.select_columns(key)[data_ids][key]: 9.562710228981814

[alexander-soare]

Maybe useful snippets:

In XarmDataset

def __getitem__(self, idx):
    start = time.perf_counter()
    item = self.hf_dataset[idx]

    if self.delta_timestamps is not None:
        start_delta = time.perf_counter()
        item, select_cols_time = load_previous_and_future_frames(
            item,
            self.hf_dataset,
            self.episode_data_index,
            self.delta_timestamps,
            tol=1 / self.fps - 1e-4,  # 1e-4 to account for possible numerical error
        )
        self._delta_time += time.perf_counter() - start_delta

    self._select_cols_time += select_cols_time
    if self.transform is not None:
        item = self.transform(item)
    self._time += time.perf_counter() - start
    return item

You have to make load_previous_and_future_frames return the select_cols_time:

start = time.perf_counter()
item[key] = hf_dataset.select_columns(key)[data_ids][key]
select_cols_time += time.perf_counter() - start

[alexander-soare]

FYI, some things I tried:

I tried doing slicing instead of advanced indexing but it didn't seem to help:

if len(torch.unique(data_ids)) == len(data_ids):
    item[key] = hf_dataset.select_columns(key)[data_ids.min().item():data_ids.max().item() + 1][key]
else:
    item[key] = hf_dataset.select_columns(key)[data_ids][key]

I tried using select, and that didn't help either: item[key] = hf_dataset.select_columns(key).select(data_ids)[key]

I tried using the keep_in_memory option when loading the dataset and that didn't help.

[alexander-soare]

In a more recent benchmark I found that even with the video datasets, more time is taken up accessing the hf_dataset, than decoding videos:

from time import perf_counter

from tqdm import trange

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
from lerobot.common.datasets.utils import load_previous_and_future_frames
from lerobot.common.datasets.video_utils import load_from_videos

fps = 15
horizon = 5

delta_timestamps = {
    "observation.image": [i / fps for i in range(horizon + 1)],
    "observation.state": [i / fps for i in range(horizon + 1)],
    "action": [i / fps for i in range(horizon)],
    "next.reward": [i / fps for i in range(horizon)],
}

dataset = LeRobotDataset("lerobot/xarm_lift_medium", delta_timestamps=delta_timestamps)

get_item_time = 0
delta_timestamps_time = 0
load_video_time = 0
transform_time = 0

def getitem(dataset: LeRobotDataset, idx: int) -> dict:
    global get_item_time, delta_timestamps_time, load_video_time, transform_time

    start = perf_counter()
    item = dataset.hf_dataset[idx]
    print("Get item:", (get_item_time := get_item_time + perf_counter() - start))

    start = perf_counter()
    if dataset.delta_timestamps is not None:
        item = load_previous_and_future_frames(
            item,
            dataset.hf_dataset,
            dataset.episode_data_index,
            dataset.delta_timestamps,
            dataset.tolerance_s,
        )
    print("Delta timestamps:", (delta_timestamps_time := delta_timestamps_time + perf_counter() - start))
    

    start = perf_counter()
    if dataset.video:
        item = load_from_videos(
            item,
            dataset.video_frame_keys,
            dataset.videos_dir,
            dataset.tolerance_s,
        )
    print("Load video:", (load_video_time := load_video_time + perf_counter() - start))

    start = perf_counter()
    if dataset.transform is not None:
        item = dataset.transform(item)
    print("Transform:", (transform_time := transform_time + perf_counter() - start))

    print()
    return item


for idx in trange(len(dataset)):
    getitem(dataset, idx)

The last iteration prints:

Get item: 3.056344897253439
Delta timestamps: 88.46708550985204
Load video: 65.91246535727987
Transform: 0.009780120220966637

[aliberts]

@lhoestq do you have any insights on this issue?

[alexander-soare]

I've been able to pin this down partially to the fact that we use an outer loop over keys to load the delta timestamps (see here):

Here, I show that we can get a 2x speedup by only slicing once (under these specific conditions):

import time

from datasets import load_dataset
from tqdm import trange

N = 256
M = 10
DATASET_REPO_ID = "lerobot/xarm_push_medium_image"

chunk_cols = ["observation.state", "observation.image", "action", "next.reward"]

hf_dataset = load_dataset(DATASET_REPO_ID, split="train")
indices = range(M)


def load_chunk(hf_dataset):
    return hf_dataset.select_columns(chunk_cols)[indices]


def load_chunk_with_outer_loop(hf_dataset):
    for col_name in chunk_cols:
        hf_dataset.select_columns(col_name)[indices]


start = time.perf_counter()
for _ in trange(N):
    load_chunk(hf_dataset)
print(time.perf_counter() - start)


start = time.perf_counter()
for _ in trange(N):
    load_chunk_with_outer_loop(hf_dataset)
print(time.perf_counter() - start)


start = time.perf_counter()
for ix in trange(N * M):
    item = hf_dataset[ix]
print(time.perf_counter() - start)

Prints
0.510965118999593
0.9346720789999381
0.4454371000001629

[Joeclinton1]

Oh, cool, I think I actually fixed this! Made delta timesteps access on the hf_dataset from the Arrow dataset 10x faster. Just wanted to flag it since it seems like my solution might have been missed. I shared details in the #datasets channel on Discord—happy to write up more + open a PR if there's interest! Just didn’t get any replies there, so wasn’t sure if it was worth pushing further.

[rokas-bendikas]

@Joeclinton1 Hey, is your solution merged already and if not, do you mind sharing it here? I Have an issue where I'm using delta_timestamps for multiple cameras and the performance is extremely slow. Example of my code:

deltas = [t / ds_meta.fps for t in range(cfg.history_length + 1)]
dataset = LeRobotDataset(
cfg.dataset_repo_id,
cfg.dataset_root,
delta_timestamps={
"observation.images.rgb.randomized_0": deltas,
"observation.images.rgb.randomized_1": deltas,
"observation.images.rgb.randomized_2": deltas,
"observation.images.rgb.randomized_3": deltas,
"observation.images.rgb.randomized_4": deltas,
},
)

I found that the bottleneck is:

def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
return {
key: torch.stack(self.hf_dataset.select(q_idx)[key])
for key, q_idx in query_indices.items()
if key not in self.meta.video_keys
}

Which runs this function for each batch item.

[lhoestq]

btw using datasets>=4 you can do hf_dataset[key][q_idx] which is fast (maybe it can help here, but not sure what the bottleneck is)

can you check what step in this _query_hf_dataset function is the actual bottleneck ?

[Joeclinton1]

@rokas-bendikas @lhoestq So I solved this problem around 6 months ago, and I talked about it on discord and asked if they would like me to make a pr, but no one responded. This commit has the solution: 6bee859 I am happy to rebase these changes on a latest version and make a pr but I was waiting for someone on the team to say that's something they'd like.

[rokas-bendikas]

@lhoestq it's the self.hf_dataset.select(q_idx)[key] that is the bottleneck. As suggested I used direct indexing, works much faster. Solution (untested, use at your own discretion):

def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
    posted_data = {}
    for key, indices in query_indices.items():
        if key not in self.meta.video_keys:
            posted_data[key] = torch.stack(
                [self.hf_dataset[idx][key] for idx in indices]
            )
    return posted_data