Concepts / Writing a Custom Collate Function for Batching

Writing a Custom Collate Function for Batching

A partial function is a new version of an existing function with some arguments prefilled, created with partial from Python's functools standard library.

  • Programming

Why Batching Needs Extra Settings

When you fine-tune a model, the PyTorch DataLoader class does the work of gathering examples into batches. The DataLoader is the component that groups individual examples together before they reach the model. You can customize that batching step by supplying your own collate function. In this part of the chapter, that function is named custom_collate_fn.

The collate function needs to know two settings that were chosen earlier: the device to send tensors to, and the maximum length a sequence may have. Those values are not part of the data itself. They are settings about how the data should be handled, so they have to reach the function some other way. This article explains the technique the source uses to get them there.

The solution is to bake those settings into a new callable ahead of time. The DataLoader then only ever sees one function that already carries the settings. Before reading on, try to hold the problem in mind: the function needs two extra values, and the DataLoader should not have to supply them.

Predicting What a Partial Does

Before looking at the mechanism, make a prediction. Suppose you have one function that takes a value and a device keyword, and you make a second name for it that always fills in the device keyword. Does the new name still require you to supply the device each time you call it? And does the original function change?

What do you think happens?

After you create a second name that always fills in the device keyword, does the original function change, and does the new name still need the device each time?

  • The original function changes, and the new name does not need the device

  • The original function is unchanged, and the new name does not need the device

  • The original function is unchanged, and the new name still needs the device

  • The original function runs immediately when the new name is created

Reveal answer
Answer:

The original function is unchanged, and the new name does not need the device each time.

partial does not rewrite the function and does not run it. It returns a new callable that remembers the arguments you gave it. Calling that new callable with the remaining arguments includes the remembered values automatically, so the original function keeps its own signature.

How partial Creates a New Callable

Definition

A partial function is a new version of an existing function with some arguments prefilled, created with partial from Python's functools standard library.

The key idea is that partial does not rewrite the function and does not run it. It returns a new callable that remembers the arguments you gave it. When you later call that callable with the remaining arguments, the remembered values are included in the call. This is why the technique is useful for reuse: you define the settings once, and every later use of the new name inherits them.

devicestored keyword valuebatchsupplied at call timeallowed_max_lengthstored as 1024
When the partial object is called with only the batch, which values come from the partial's stored settings and which come from the call?

A Small Partial Example

A short example with an ordinary function makes the behavior easier to see. The function below is not part of the batching problem. It only illustrates how partial stores a keyword value and how the original stays available.

Code walkthrough
from functools import partial

def power(base, exponent):
    return base ** exponent

square = partial(power, exponent=2)

print(power(3, 3))
print(square(5))
print(power(3, 3))
  1. 1
    The import line brings partial into scope.
  2. 2
    The power function is defined normally and is not modified.
  3. 3
    partial(power, exponent=2) returns a new callable named square. It does not call power.
  4. 4
    print(power(3, 3)) calls the original function with both arguments and prints 27.
  5. 5
    print(square(5)) calls the partial with only base. The remembered exponent=2 is included, so the result is 5 to the power 2, which is 25.
  6. 6
    print(power(3, 3)) prints 27 again, showing that the original function was not changed by creating the partial.
  • partial is imported from functools
  • power is the original function with two parameters
  • square stores exponent=2 and returns a new callable
  • Calling square(5) supplies only base, and the stored exponent is included
  • The original power still needs both arguments and still returns 27 for power(3, 3)
Output
27
25
27

Building customized_collate_fn

The source builds the partial function in two steps. The first step is to import partial from functools. The second step is to assign a new name to the result of calling partial on three things: the original function custom_collate_fn, the device value, and the number 1024 assigned to allowed_max_length.

Code walkthrough
from functools import partial

customized_collate_fn = partial(custom_collate_fn, device=device, allowed_max_length=1024)
  1. 1
    The first line imports the partial function from Python's functools module.
  2. 2
    The second line calls partial with the original function custom_collate_fn as its first argument.
  3. 3
    The device value is passed by keyword name, so it is clear which parameter it fills.
  4. 4
    The number 1024 is passed by the keyword allowed_max_length, so it fills that parameter.
  5. 5
    The result is assigned to customized_collate_fn, which is a different name from custom_collate_fn. The two names refer to different callables.
  • partial comes from functools
  • custom_collate_fn is the original function
  • device=device fills the device parameter by keyword
  • allowed_max_length=1024 fills the length parameter by keyword
  • customized_collate_fn is the new name for the partial

Tracing a Batch Through the Partial

Follow what happens when the DataLoader needs a batch. The DataLoader calls customized_collate_fn with the batch of samples. The partial object supplies the stored device and allowed_max_length values and forwards the batch together with them to the original custom_collate_fn. The DataLoader never handles those two settings itself.

DataLoadergathers examples intobatchescustomized_collate_fnpartial objectdevicestored valuecustom_collate_fnoriginal functionallowed_max_length =1024stored valuecalls with batchholdsholdsforwards batch and stored values
When the DataLoader fetches a batch, the call passes through the partial object, which adds the stored device and allowed_max_length values before reaching the original collate function.

Read the diagram from top to bottom. The DataLoader's call enters the partial object. The partial object holds the two prefilled values, and it passes them along with the batch to the original function. The original function receives all of its arguments, even though the DataLoader supplied only one.

The Original Function Stays Unchanged

Creating customized_collate_fn leaves custom_collate_fn as it was. Calls made directly to the original name still need every argument that the original function requires. Calls made through the new name need only the remaining argument, because the stored values are attached automatically.

custom_collate_fnbatch, device,allowed_max_lengthcustom_collate_fnbatch, device,allowed_max_lengthcustomized_collate_fnbatch only
Creating the partial adds a second name with fewer required arguments, while the original function's signature is left as it was.

Plugging the Partial into the DataLoader

The last step in this part of the chapter is to set up the data loaders as was done previously, with one difference: the custom collate function now handles batching. The partial version, customized_collate_fn, is the function that plays this role.

The DataLoader does not need to know about the device or the length limit, because those values were already attached to the function. From the DataLoader's point of view, it is handing batches to one ordinary collate function. The settings are hidden inside the partial, which keeps the data loader setup the same as before.

Mistakes When Using a Partial Collate Function

  • Common mistake

    Expecting the original name to pick up the prefilled values

    What it looks like

    After creating customized_collate_fn, calling custom_collate_fn with only the batch.

    Why it fails

    The partial creates a new callable. It does not change the original function's signature.

    Do this instead

    Call custom_collate_fn with every argument it needs, or call customized_collate_fn when you want the stored values included.

  • Common mistake

    Assuming partial runs the function right away

    What it looks like

    Writing customized_collate_fn = partial(custom_collate_fn, device=device, allowed_max_length=1024) and expecting batches to be collated at that line.

    Why it fails

    partial does not run the function. It returns a new callable that waits to be called.

    Do this instead

    Expect the function to run only when the DataLoader or your code later calls customized_collate_fn.

  • Common mistake

    Changing the 1024 value without considering the model

    Why it fails

    The source ties the value to the maximum context length supported by GPT-2, so the value is chosen to match the model.

    Do this instead

    Keep allowed_max_length aligned with the model's supported context length.

  • Common mistake

    Passing the device and length to the DataLoader as well

    Why it fails

    Those values are already attached to the partial, so the DataLoader does not need them.

    Do this instead

    Pass customized_collate_fn as the collate function and leave the other DataLoader settings as they were.

Practice Predicting Partial Calls

MEDIUM

Using the power and square example from earlier, predict the output of these three lines before you run them: print(square(4)), print(power(2, 5)), and print(square(4, 3)). Explain which of the three would fail and why.

Hints
  • Ask which arguments are stored inside square and which must be supplied at call time.
  • Think about whether the original power still requires both arguments.
  • Consider what happens when you pass a third positional argument to a function that only has two parameters, with one of them already filled by a keyword.

Summary of the Partial Collate Pattern

  1. 1
    A partial function is a new version of an existing function with some arguments prefilled, created with partial from functools.
  2. 2
    partial does not rewrite or run the original function. It returns a callable that remembers the arguments you gave it.
  3. 3
    customized_collate_fn = partial(custom_collate_fn, device=device, allowed_max_length=1024) bakes in the device and the 1024 length limit.
  4. 4
    The DataLoader uses customized_collate_fn as its custom collate function, so it never needs to supply the device or length itself.
  5. 5
    The original custom_collate_fn keeps its own signature and still needs all of its arguments when called directly.

At a glance

  • A partial function is a new version of an existing function with some arguments prefilled, created with partial from functools.
  • partial returns a new callable that remembers the given arguments and does not run or rewrite the original function.
  • The partial customized_collate_fn stores device and allowed_max_length=1024 so the DataLoader can use it as a plain collate function.
  • The original custom_collate_fn keeps its own signature and still needs all of its arguments when called directly.

Related Topics

Related concepts

More concepts starting with W