Python Sentence Transformers Batch Encoding and GPU

When you need embeddings for thousands of sentences, calling model.encode() on each sentence individually is wasteful. The sentence-transformers library is built around the idea of encoding lists of texts in one call, and on a GPU that single call can process many sentences in parallel. This article explains how python sentence transformers batch encoding and gpu work together, how to control the batch size, and what to watch out for when scaling up.

Why Batch Encoding Matters for Sentence Transformers on GPU

A transformer model processes a fixed-length sequence of tokens. When you call encode() with a single sentence, the GPU executes the forward pass for that one sequence. With a batch, the GPU processes multiple sequences simultaneously, sharing the model weights and amortizing the overhead of kernel launches and memory transfers. The result is higher throughput: more sentences per second, and often lower total latency for a large corpus.

GPU hardware is designed for parallel computation. A batch of 32 sentences can be processed in nearly the same time as a single sentence, because the GPU has many cores that operate on different rows of the batch matrix. This is the core reason to use batch encoding rather than a loop of single-sentence calls.

Basic Batch Encoding with sentence-transformers

The SentenceTransformer class provides an encode() method that accepts a list of strings. Here is the simplest form:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('all-MiniLM-L6-v2')
sentences = [
    "The cat sits on the mat.",
    "A dog runs through the park.",
    "Machine learning models require data."
]

embeddings = model.encode(sentences)

The embeddings variable is a NumPy array of shape (3, 384) for this model. Internally, the library tokenizes all sentences, pads them to the same length, and runs the transformer on the whole batch. The default batch size is 32, but you can change it.

Controlling Batch Size and Device

The encode() method accepts a batch_size parameter. You can also specify the device explicitly, for example device='cuda' to force GPU usage. If you have a GPU and PyTorch detects it, the model is automatically placed on the GPU, but you can be explicit:

embeddings = model.encode(sentences, batch_size=64, device='cuda')

When you pass a batch size larger than the number of sentences, the library processes all sentences in one batch. When the batch size is smaller, it splits the list into chunks and processes each chunk sequentially. The show_progress_bar parameter can help you monitor progress for large corpora:

embeddings = model.encode(sentences, batch_size=128, show_progress_bar=True)

GPU Memory and Choosing the Right Batch Size

GPU memory is finite. Each sentence in a batch consumes memory for its token embeddings, attention masks, and intermediate activations during the forward pass. The exact amount depends on the model size and the maximum sequence length. If the batch size is too large, you will get a CUDA out-of-memory error. The error typically looks like:

RuntimeError: CUDA out of memory.

To avoid this, you need to find the largest batch size that fits on your GPU. A practical approach is to start with a small value, say 16, and double it until you hit an error, then back off. The optimal batch size also depends on the sequence length of your sentences. Longer sentences consume more memory because the attention matrix grows quadratically with sequence length.

This table is illustrative. The actual values depend on the model and GPU. The key is to monitor memory usage and choose the largest batch that does not cause an OOM error.

Normalization and Output Formats

The encode() method has several useful options. By default, it returns a NumPy array of float32 vectors. If you need PyTorch tensors, set convert_to_tensor=True. If you need normalized vectors for cosine similarity, set normalize_embeddings=True. This is often done for retrieval tasks:

embeddings = model.encode(sentences, normalize_embeddings=True, convert_to_tensor=True)

Normalization is performed on the GPU, so it does not add a significant overhead. It is better to normalize the entire batch at once rather than doing it manually in a loop.

Common Failure Modes and How to Handle Them

One common mistake is mixing CPU and GPU tensors. If you load the model on GPU but then pass a list of strings, the library handles the transfer automatically. However, if you try to manually move inputs, you can run into device mismatch errors. Always let the library manage device placement unless you have a specific reason not to.

Another issue is using a batch size that is too large for the GPU. The library does not automatically reduce the batch size; it simply raises an error. You can catch this error and retry with a smaller batch:

try:
    embeddings = model.encode(sentences, batch_size=256)
except RuntimeError as e:
    if 'out of memory' in str(e):
        embeddings = model.encode(sentences, batch_size=64)
    else:
        raise

This is a simple fallback, but in production you should pre-determine a safe batch size based on your GPU and typical sequence length.

When to Use Batch Encoding vs. Streaming or Incremental Encoding

For a one-off script that encodes a few thousand sentences, a single encode() call with a reasonable batch size is sufficient. For very large corpora that do not fit in memory, you might need to process the data in chunks. The library does not provide a built-in streaming API, but you can read your data in chunks and call encode() on each chunk, reusing the same model. This avoids loading all sentences into RAM at once.

Batch encoding is also useful when you need to compare many sentences against each other, such as in clustering or semantic search. The embeddings are produced in a single call, and you can then compute similarity matrices using NumPy or PyTorch operations.

If you are encoding a single sentence repeatedly, for example in an online inference service, batch encoding is not applicable. In that case, you should keep the model on GPU and call encode() with a single sentence, but the overhead of a single forward pass is still small. The benefit of batch encoding becomes significant when you have at least a few dozen sentences to process at once.

A final consideration is the interaction with torch.no_grad(). The encode() method already disables gradient computation internally, so you do not need to wrap it yourself. This saves memory and speeds up inference.

When you move to a multi-GPU setup, sentence-transformers supports data parallelism through PyTorch. You can pass device='cuda:0' or use torch.nn.DataParallel manually, but the library's encode() method does not automatically split the batch across multiple GPUs. For most applications, a single GPU with a well-chosen batch size is sufficient. If you need to scale further, consider using the sentence-transformers training API or a dedicated inference server.