Comfortable pin_memory option

Hi there,
I was looking through the framework to find out how to enable the pin_memory option in pytorch dataloaders. Unfortunately I have not found a way to do so using the data block api or am I missing something?

I am not sure if the trivial solution
data.train_dl.pin_memory = True
works since this value is only checked in the DataLoaderIter constructor and I have no clue when this is constructed (but I believe it’s somewhere before the end of the data block pipeline)

The only option (without modifying my fastai installation) to pass the argument in the DataLoader constructor that I see is to replace the
intercept_args(...)
function from basic_data.py and do
torch.utils.data.DataLoader.__init__ = modified_intercept_args

So did I just miss the place in the data blocks api where I can define it or is there any good reason why I cannot specify this? I could also not find any functions in the documentation stating they would pass their **kwargs to the DataLoader (I found this for some other **kwargs).

I’ll gladly make a PR if you think it’s a good idea to include this. (.databunch(pin_memory=True) seems like a sensible place?)

EDIT: Here’s the PR: https://github.com/fastai/fastai/pull/1674

Greetings,
Dome

1 Like

The intended way for now is data.train_dl = data.train_dl.new(pin_memory=True).

Otherwise, we could have **dl_kwargs in DataBunch.create (which is where .databunch end up) that would be passed to the pytorh DataLoader call. If you want to make a PR with this, it would need to be included in DataBunch.create, TextClasDataBunch.create and maybe TextLMDataBunch.create.