Running on a webdataset

[nicolas-dufour]

Hi,
Is there a simple way to run this code on a webdataset?

Thanks!

[huyvvo]

Hello,

This code receives as input embeddings stored in a single .npy file so how the original dataset is represented is not relevant. To run the code, you can extract embeddings for all images in the webdataset and save them in an .npy file.

Hope this answers your question!

[nicolas-dufour]

Thanks!
But how do you deal with memory issues then if the npy don't fit into memory for big datasets

[vkhalidov]

how do you deal with memory issues then if the npy don't fit into memory for big datasets

Please note that mmap_mode="r" is used when loading the embeddings (see scripts/run_distributed_kmeans.py#L51). So data loading is fast and RAM footprint is low.

Full data for a given worker gets loaded in src/distributed_kmeans_gpu.py#L96 and it's more a bottleneck for GPU memory than for RAM. However, this is parallelized across workers. So distributing sufficiently across GPUs allows to avoid GPU OOM.

[huyvvo]

Thanks! But how do you deal with memory issues then if the npy don't fit into memory for big datasets

For saving embeddings into a big .npy file, you can use "numpy.lib.format.open_memmap" funtionality. Another solution is to save embeddings into multiple files and use a class that imitate numpy.array, then replace line scripts/run_distributed_kmeans.py#L51 with the class' constructor.