How does nlp.pipe work to prevent OOM?

[BramVanroy]

I am using a transformer-based model with require_gpu. I can confirm that my code works even when batch size is set to 1024 in pipe() on a 16GB GPU. So I wonder, does pipe internally do something special (smaller minibatches based on some heuristic or catching OOM and systematically lowering batch size?) to ensure that out-of-memory issues don't occur?

[svlandeg]

Hi Bram! Happy to hear things are working well for you.

You can find the code for nlp.pipe here: https://github.com/explosion/spaCy/blob/master/spacy/language.py#L1487. You can see that batch_size is used as such. If n_process>1, the function self._multiprocessing_pipe is called which in turn calls util.minibatch with the same batch_size.

Is there a particular reason why you're asking? (I might not have understood the question)

[BramVanroy]

Hi Sofie, thanks for the reply. I was making a mistake. I assumed that pipe processed the given Iterable as a single batch. So I was creating mini-batches of a given size before feeding each batch to pipe(). Looking at the code, I just now realized that pipe() uses its own batch_size argument, which is None by default. If it is None, it will use the default batch size of the Language instance. This is evident in the following example.

nlp = spacy.load("nl_udv25_dutchalpino_trf")
nlp.batch_size
# 64
nlp2 = spacy.load("en_core_web_sm")
nlp2.batch_size
# 256

The transformer-based model has a smaller batch size by default, which makes sense.

So, as you said, pipe does not process the given texts as a single batch, but chunks them into batch_size-sized minibatches. I was doing the former and wondering how pipe then seemed to work for any batch size. But now it is obvious, it simply chunked the given input into smaller batches.

So you do not have to do something like this:

for batch in minibatch(lines, size=batch_size):
    docs = nlp.pipe(batch)

but instead

docs = nlp.pipe(lines, batch_size=batch_size)

[svlandeg]

Ah, right. That makes sense. Happy you figured it out :-)