Running out of memory during dataset.map() with `AutoFeatureExtractor.from_pretrained("facebook/hubert-large-ls960-ft")`
Hi, I’m having an issue of running out of memory when trying to use themapfunction on a Dataset. A simplified, (mostly) reproducible example (on a 16 GB RAM) is below.Let’s say I have a dataset of 1000 audio files of varying lengths from 5 seconds to 20 seconds, all sampled in 16 kHz. Assume I have the followingDatasetobject to represent that:import numpy as np
from datasets import Dataset

dataset_size = 1000

dummy_data = {
    "audio_file": [
        {
            'path': "ignore-me.wav",
            'array': np.random.random(
                np.random.choice(np.arange(5* 16_000, 20 * 16_000), 1)[0].astype(int)  
                # an audio file of a length between 5s and 20s
            ).astype(np.float32),
            'sampling_rate': 16000
        } for _ in range(dataset_size) 
    ],
    "labels": [np.random.choice(['A', 'B']) for _ in range(dataset_size)],
}
ds = Dataset.from_dict(dummy_data)
del dummy_dataI want to use the"facebook/hubert-large-ls960-ft"feature extractor to preprocess this data set to train a Hubert model. This is my preprocessing code.from transformers import AutoFeatureExtractor

feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/hubert-large-ls960-ft")
def preprocess_function(examples):
    audio_arrays = [x['array'] for x in examples['audio_file']]
    inputs = feature_extractor(
        audio_arrays,
        sampling_rate=feature_extractor.sampling_rate,
        max_length=160_000,
        truncation=True,
        padding=True
    )
    return inputsWhen I apply the following, it causes my machine with 16 GB RAM to crash (notebook crashes).htopsuggests that the process is exhausting all the memory.ds.map(
    preprocess_function,
    remove_columns=['audio_file'],
    batched=True)Question:I believe that this behaviour does not align withThe magic of memory mapping. Most likely, I’m doing something wrong here. How can I fix it?PS:In my actual data set, there are audio clips from 5s to several minutes long and I have about 15,000 total audio clips. I believe it’s quite a small data set, considering the size of other speech data sets available.I obtain the variabledsabove on the actual data set by using thefrom_pandasmethod and applyingdatasets.Audiofunction to load as an audio clip.Versions:transformers=4.18.0,datasets=2.1.0Thanks!
2022-06-08T00:39:11Z
Hey there, I just ran into this issue when processing images, and found a potential solution in the docs- maybe it will work for you.Inthis sectionof the docs, it says:Dataset.map()takes up some memory, but you can reduce its memory requirements with the following parameters:batch_sizedetermines the number of examples that are processed in one call to the transform function.writer_batch_sizedetermines the number of processed examples that are kept in memory before they are stored away.I wasn’t usingbatched=True, so I just had to changewriter_batch_sizeto something small (in my case I set it to 10). In your case, you might want to play with changingbatch_sizeas well.Edit: actually I just found that thisdoes notsolve my problem, I was mistaken.@lhoestq@mariosaskoany tips here?
2022-06-08T01:14:09Z
@naterawThis is actually a duplicate question. See the solution tothis questionfor the fix.Spoiler: you have to passcache_file_namein.map().
2022-06-08T02:33:54Z
Oh wow, interesting! I was using pandas too to create the dataset originally, so it was the same exact problem. Thanks for the linkThis should solve my issue too