Define what the loader must deliver
Write the output contract before optimizing: number of elements, type of each field, dimensions, target range, and the rule for incomplete inputs. Distinguish the example identifier from its position in a batch. A transformation can change a shape or filter an input; the training program must know whether that is allowed.
Take a representative sample including an ordinary file, an edge case, and the last element of the dataset. Open each element with exactly the same preparation as the Dataset. Then examine their collation. A successful individual access does not prove that multiple results can be stacked. For text, document padding and mask; for an image, channels, dimensions, and axis order.
Set the scope: local or remote data, decoding included or not, fixed or random transformations. Keep it consistent between two runs; an apparent gain can come from work that was removed.
Go back to a single process to read the error
First reproduce with num_workers=0, shuffle=False, and a small batch. Loading then runs in the main process and the error trace is usually more readable. The DataLoader documentation recommends this option for debugging. Log the identifier of the element that fails before it is decoded, without copying its sensitive content into the logs.
Proceed by separation: raw access, transformation, collate_fn, then transfer. If the path fails before the transfer, changing CUDA is not the first thing to try. If it only hangs with multiple workers, examine the objects and resources passed to those processes. Compare the first iteration and the following ones: worker startup can explain an initial wait without establishing a recurring problem.
A timeout can make a wait visible, but it does not fix an unavailable source or a stuck worker. Keep the last known step and reduce the number of inputs instead of increasing that timeout indefinitely.
Worked example: three expected channels, a different image
Consider four educational records. The first three yield a tensor of shape [3, 16, 16], the fourth [1, 16, 16]. With a contract requiring three channels, the fourth element must be identified before collation. This scenario was not executed here; it describes an expected result based on the chosen shapes.
The function below assumes that each record has the fields id, x, and y, that x is a CPU tensor, and that y is an integer index. It rejects the inconsistency instead of silently discarding the image. For your project, decide explicitly whether a monochrome image should be converted to three channels or rejected at import. That decision depends on the meaning of the data and the preprocessing expected by the model.
After the fix, all four identifiers must remain present and the collated tensor must have shape [4, 3, 16, 16]. Add a guard suited to the targets: a correctly sized image can still carry an invalid annotation.
import torch
from torch.utils.data import DataLoader
def assemble(records):
for item in records:
if tuple(item["x"].shape) != (3, 16, 16):
raise ValueError(f"Unexpected shape for {item['id']}")
return {
"ids": [item["id"] for item in records],
"x": torch.stack([item["x"] for item in records]),
"y": torch.tensor([item["y"] for item in records],
dtype=torch.long),
}
# dataset is your Dataset producing the records described.
# In a multiprocess script, create the loader under the main guard.
if __name__ == "__main__":
loader = DataLoader(dataset, batch_size=4, num_workers=0,
shuffle=False, collate_fn=assemble)
iterator = iter(loader)
batch = next(iterator)Reintroducing workers without changing the data
Go from zero to a small number of workers while keeping the batch, order and transformations. Test a full epoch, then a second one: some errors only appear when an iterator restarts or after resources have been consumed. Increasing parallelism is only useful if the preparation work can actually proceed in parallel.
The start methods depend on the system and the Python version. With spawn, protect the program entry point with if __name__ == '__main__' and define Dataset, collate_fn and worker functions at module level rather than in local lambdas. The process documentation also explains why inherited locks or threads can cause deadlocks. Keep the initialization of per-process accesses when the library requires it.
For an IterableDataset, check the partition between workers using identifiers: several workers must not each consume the entire same stream. Do not judge only the number of batches; also look for duplicates and missing elements.
Measuring wait time and throughput with a clear unit
Use two complementary observations. A pass over the loader alone counts the examples prepared during a defined interval. An integrated pass examines what happens when the model consumes that data. The first helps isolate preparation; it does not automatically represent training throughput.
In your protocol, count examples actually delivered, then divide by elapsed seconds. State the passes excluded for startup, the data cache, the transformations and the number of repetitions. Keep the values from each pass instead of selecting only the best. The table below is a log sheet: no performance is filled in.
If the shapes vary, an examples-per-second figure can mask a change in load. Add the relevant unit, such as pixels decoded or tokens actually prepared, while also keeping the examples. To locate the waits in the full loop, name the reading of the next batch separately from the computation.
Scroll the table to read all columns.| Setting | Items checked | Observed duration | Expected conclusion |
|---|---|---|---|
| workers=0 | Identifiers, shapes, targets | To be measured in seconds | Correct baseline |
| Small number of workers | Same set of inputs | To be measured in seconds | Real gain or overhead |
| Same setting, second epoch | No loss or duplication | To be measured in seconds | Effect of startup and caches |
Handling memory, prefetching and transfers separately
Workers and queued batches consume host memory. Monitor it during your test before concluding that only VRAM matters. Deeper prefetching can shift the wait while increasing usage; it does not guarantee more results per second. First reduce the suspect variable and compare the same scope.
pin_memory and non-blocking transfers concern the movement of data to an accelerator. The PyTorch optimization recipe presents them as levers to examine together with the hardware and the workload. They do not fix incorrect decoding. Start with CPU data in the workers, then organize the transfer in the process driving the computation. The benefit and the effective overlap must be observed, not assumed.
If persistent_workers is used, keep in mind the resources and state retained between epochs. A setting that looks fine on a single batch is not enough to verify that files are closed or that the data source is refreshed.
Accept a setting only if the data remains correct
The expected result is a loop that receives all the intended inputs, within the chosen scope, without silent errors. Compare identifiers and targets before and after optimization. Explain drop_last if you discard the last incomplete batch. If the transforms are random, check their policy rather than demanding pixel equality, which would contradict that policy.
Keep the simplest setting that meets the measured need. Increasing the number of workers may not improve anything if storage, decoding, or the model is already the bottleneck. A server's CPU, RAM, and storage resources cannot be inferred from its GPU name: specify these requirements separately when you prepare your Kernodeck environment.