Datasets: Creating dataset consumes too much memory

[AmitMY]

Moving this issue from https://github.com/huggingface/datasets/pull/722 here, because it seems like a general issue.

Given the following dataset example, where each example saves a sequence of 260x210x3 images (max length 400):

    def _generate_examples(self, base_path, split):
        """ Yields examples. """

        filepath = os.path.join(base_path, "annotations", "manual", "PHOENIX-2014-T." + split + ".corpus.csv")
        images_path = os.path.join(base_path, "features", "fullFrame-210x260px", split)

        with open(filepath, "r", encoding="utf-8") as f:
            data = csv.DictReader(f, delimiter="|", quoting=csv.QUOTE_NONE)
            for row in data:
                frames_path = os.path.join(images_path, row["video"])[:-7]
                np_frames = []
                for frame_name in os.listdir(frames_path):
                    frame_path = os.path.join(frames_path, frame_name)
                    im = Image.open(frame_path)
                    np_frames.append(np.asarray(im))
                    im.close()

                yield row["name"], {"video": np_frames}

The dataset creation process goes out of memory on a machine with 500GB RAM.
I was under the impression that the "generator" here is exactly for that, to avoid memory constraints.

However, even if you want the entire dataset in memory, it would be in the worst case
260x210x3 x 400 max length x 7000 samples in bytes (uint8) = 458.64 gigabytes
So I'm not sure why it's taking more than 500GB.

And the dataset creation fails after 170 examples on a machine with 120gb RAM, and after 672 examples on a machine with 500GB RAM.

---

Info that might help:

Iterating over examples is extremely slow.

If I perform this iteration in my own, custom loop (Without saving to file), it runs at 8-9 examples/sec

And you can see at this state it is using 94% of the memory:

And it is only using one CPU core, which is probably why it's so slow:

[lhoestq]

Thanks for reporting.
In theory since the dataset script is just made to yield examples to write them into an arrow file, it's not supposed to create memory issues.

Could you please try to run this exact same loop in a separate script to see if it's not an issue with PIL ?
You can just copy paste what's inside _generate_examples and remove all the code for datasets (remove yield).

If the RAM usage stays low after 600 examples it means that it comes from some sort of memory leak in the library, or with pyarrow.

[AmitMY]

Here's an equivalent loading code:

images_path = "PHOENIX-2014-T-release-v3/PHOENIX-2014-T/features/fullFrame-210x260px/train"

for dir_path in tqdm(os.listdir(images_path)):
    frames_path = os.path.join(images_path, dir_path)
    np_frames = []
    for frame_name in os.listdir(frames_path):
        frame_path = os.path.join(frames_path, frame_name)
        im = Image.open(frame_path)
        np_frames.append(np.asarray(im))
        im.close()

The process takes 0.3% of memory, even after 1000 examples on the small machine with 120GB RAM.

I guess something in the datasets library doesn't release the reference to the objects I'm yielding, but no idea how to test for this

[lhoestq]

I've had similar issues with Arrow once. I'll investigate...

For now maybe we can simply use the images paths in the dataset you want to add. I don't expect to fix this memory issue until 1-2 weeks unfortunately. Then we can just update the dataset with the images. What do you think ?

[AmitMY]

If it's just 1-2 weeks, I think it's best if we wait. I don't think it is very urgent to add it, and it will be much more useful with the images loaded rather than not (the images are low resolution, and thus papers using this dataset actually fit the entire video into memory anyway)

I'll keep working on other datasets in the meanwhile :)

[lhoestq]

Ok found the issue. This is because the batch size used by the writer is set to 10 000 elements by default so it would load your full dataset in memory (the writer has a buffer that flushes only after each batch). Moreover to write in Apache Arrow we have to use python objects so what's stored inside the ArrowWriter's buffer is actually python integers (32 bits).

Lowering the batch size to 10 should do the job.

I will add a flag to the DatasetBuilder class of dataset scripts, so that we can customize the batch size.

[AmitMY]

Thanks, that's awesome you managed to find the problem.

About the 32 bits - really? there isn't a way to serialize the numpy array somehow? 32 bits would take 4 times the memory / disk space needed to store these videos.

Please let me know when the batch size is customizable and I'll try again!

[lhoestq]

The 32 bit integrers are only used in the writer's buffer because Arrow doesn't take numpy arrays correctly as input. On disk it's stored as uint8 in arrow format ;)

[AmitMY]

I don't expect to fix this memory issue until 1-2 weeks unfortunately.

Hi @lhoestq
not to rush of course, but I was wondering if you have a new timeline so I know how to plan my work around this :)

[lhoestq]

Hi ! Next week for sure :)

[lhoestq]

Alright it should be good now.
You just have to specify _writer_batch_size = 10 for example as a class attribute of the dataset builder class.

[AmitMY]

I added it, but still it consumes as much memory

https://github.com/huggingface/datasets/pull/722/files#diff-2e0d865dd4a60dedd1861d6f8c5ed281ded71508467908e1e0b1dbe7d2d420b1R66

Did I not do it correctly?

[lhoestq]

Yes you did it right.
Did you rebase to include the changes of #828 ?

EDIT: looks like you merged from master in the PR. Not sure why you still have an issue then, I will investigate

[AmitMY]

Hi @lhoestq, any update on this?
Perhaps even a direction I could try myself?

[lhoestq]

Sorry for the delay, I was busy with the dataset sprint and the incredible amount of contributions to the library ^^'

What you can try to do to find what's wrong is check at which frequency the arrow writer writes all the examples from its in-memory buffer on disk. This happens here in the code.

The idea is that write_on_file writes the examples every writer_batch_size examples and clear the buffer self. current_rows. As soon as writer_batch_size is small enough you shouldn't have memory issues in theory.

Let me know if you have questions or if I can help.

Since the dataset sprint is over and I will also be done with all the PRs soon I will be able to go back at it and take a look.

[AmitMY]

Thanks. I gave it a try and no success. I'm not sure what's happening there