DataLoader with pin_memory and num_workers

16 of december of 2024

DataLoader + pin_memory

Disclaimer: This post has been translated to English using a machine translation model. Please, let me know if you find any mistakes.

In PyTorch, when training neural networks, especially on large datasets, leveraging the DataLoader with pin_memory=True and setting num_workers to a positive number significantly increases performance.

pin_memory=True allows for faster transfer of data to the GPU by keeping it in pinned (page-locked) memory.

At the same time, num_workers determines the number of subprocesses used for data loading, allowing for asynchronous data retrieval without blocking the GPU calculation

This combination minimizes the GPU downtime, ensuring more efficient use of hardware resources and faster model training times.

data_loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

Frequently asked questions

What exactly does pin_memory=True do in PyTorch's DataLoader, and why does it speed up the transfer to the GPU?

`pin_memory=True` makes the `DataLoader` allocate each batch's tensors in page-locked ("pinned") host memory instead of regular pageable memory. This lets the host-to-GPU copy happen asynchronously and faster, because the CUDA driver can use DMA directly on that region without an extra intermediate copy into a temporary buffer.

What does num_workers do in DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True), and how does it keep the GPU from sitting idle?

`num_workers=4` tells the `DataLoader` to use 4 subprocesses to load and prepare batches in parallel, instead of doing it synchronously in the main process. That way, while the GPU is computing on the current batch, those workers are already asynchronously preparing the next one, which combined with `pin_memory=True` minimizes GPU idle time and speeds up training.

Continue reading

Gymnasia: architecture of an AI agent that runs on the phone

How the Gymnasia AI agent is built: a BYOK conversational coach with 12 local tools and a vision subagent that estimates macros from photos, both running entirely on the phone with no backend.

Deep Research with LangGraph (3/3): the Writer agent and final report

Third and final part of the Deep Research with LangGraph series. Implement the Writer agent that drafts the final report...

Deep Research with LangGraph (2/3): the multi-agent Research Supervisor

Second part of the Deep Research with LangGraph series. Build the Research Supervisor that coordinates several Researche...

Last posts -->

Have you seen these projects?

Gymnasia

LangGraph Deep Researcher

Tau

View all projects -->

>_ Available for projects

Do you have an AI project?

Let's talk.

maximofn@gmail.com

Machine Learning and AI specialist. I develop solutions with generative AI, intelligent agents and custom models.

Write me LinkedIn

Do you want to watch any talk?

Tomorrow's Agents: Deciphering the Mysteries of Planning, UX and Memory

AI agents, powered by LLMs, promise to transform applications. But are they simple executors today or future intelligent collaborators? To reach their...

Create your own Apple intelligence

Learn to create an IA system to execute efficiently on a device

Last talks -->

Do you want to improve with these tips?

Best practices building agents with Claude Code

Technical talk: skills, subagents, slash commands and MCPs in Claude Code

o1 prompt engineering

Create better prompts for o1 following an example

Memory profiler

See the memory usage of a script

Last tips -->

Use this locally

Hugging Face spaces allow us to run models with very simple demos, but what if the demo breaks? Or if the user deletes it? That's why I've created docker containers with some interesting spaces, to be able to use them locally, whatever happens. In fact, if you click on any project view button, it may take you to a space that doesn't work.