DataLoader: direct multi-index instead of coallate for batches on map-style datasets

username: jacobbuckman
cooked: <p>I’m trying to implement a batched data loader. When I implement a BatchedSampler, it gives me elements like this: [1,2,3], [4,5,6],…etc. Were I to load the data myself directly, I would want to just use multiindexing: train_set[idx] where idx=[1,2,3] or whatever. (Both are tensors, of course.) But it seems like the DataLoader insists on using a coallate function. It returns: coallate([train_set[1], train_set[2], train_set[3]).</p>
<p>Is there any way to disable this behavior? I’m using a weird custom object, so doing many individual queries is expensive, as is a coallate_fn call. And, I’ve already set it up so that we can multiselect quickly. I can see no good reason why we should select each one separately and then coallate.</p>
<p>Is there a way to do this?</p>

username: ptrblck
cooked: <p>You could try to use a <code>BatchSampler</code> as given in <a href="https://discuss.pytorch.org/t/how-to-use-batchsampler-with-getitem-dataset/78788/4">this code snippet</a>, which would pass multiple indices to the <code>__getitem__</code> and thus you could load multiple samples in this call.</p>

username: jacobbuckman
cooked: <p>Thanks, this is just what I was after.</p>
<p>A follow-up: my dataloader uses a custom data type, <code>Data</code>. <code>Data</code> has a function <code>pin_memory(self)</code>, which returns another object of custom type <code>Data</code> but in pinned memory. However, when I set <code>pin_memory=True</code> in my dataloader, suddenly my iterable yields objects of type <code>dict</code> instead of <code>Data</code>. Any idea what might be going on?</p>

username: ptrblck
cooked: <p>No, I don’t know how custom objects are treated in the <code>DataLoader</code> and might interact with pinning the memory. What’s your use case that you need to pin the memory in your custom object and then again in the <code>DataLoader</code>? Or are you trying to use the <code>DataLoader</code> to call into your custom <code>pin_memory</code> method?</p>

username: jacobbuckman
cooked: <p>Yes, the latter. Here’s a full description of what I am doing:</p>
<p>I have a custom <code>Dataset</code> that stores a large dataset in a compressed way, and I have code to efficiently decompress a minibatch onto the CPU, which I implement through <code>.__getitem__(index)</code> where index is batched. <code>.__getitem__</code> returns a custom data type, <code>Data</code>, which lets me conveniently work with many different fields, i.e. <code>data.x</code>, <code>data.y</code>, <code>data.meta</code>, <code>data.source</code>, etc., each of which is a CPU Tensor whose first dimension is minibatch-sized. <code>Data</code> also has a couple other convenience functions, including <code>.pin_memory()</code>, which returns a new instance of <code>Data</code> where all of the Tensors now live in pinned memory.</p>
<p>So, I want a DataLoader which takes a batch of indices, passes it to <code>.__getitem__</code>, takes the resulting <code>Data</code> object and calls <code>.pin_memory</code> on it, and then holds the result in a queue so it can later be fed to my training loop. The first half (<code>__getitem__</code> on batch of indices) was solved by the forum post you linked. The second issue (calls <code>.pin_memory</code> and returns the result) is what I’m struggling with now.</p>

username: ptrblck
cooked: <p>Interesting use case. I’m still unsure what’s causing the <code>dict</code> creation, but I guess that something might have changed in the object such that the <code>collate_fn</code> is now treating it differently.<br>
Could you try to create a custom <code>collate_fn</code> and return the desired type?</p>

username: jacobbuckman
cooked: <p>I’ve been using <code>collate_fn=lambda x:x</code> in an effort to avoid doing anything weird. I’ve confirmed that the input and output are both of custom type <code>Data</code>.</p>

username: jacobbuckman
cooked: <p>I traced it back to <code>_utils.pin_memory.pin_memory</code>, which uses an if/else statement to outline a bunch of options for the pinning. One of these options checks if the object is a <code>collections.abc.Mapping</code>, and if so, returns a dict with each item pinned. Since my custom object is a dict subclass, this is triggering instead of my class’s custom <code>pin_memory</code>.</p>
<p>I’ve monkey-patched the function on my own PyTorch build, so I’m good, but it might be worth fixing for the library. Maybe you could modify that function so that it exclusively triggers if the object in question is a dict.</p>

username: ptrblck
cooked: <p>If you think this behavior is wrong and could be improved, feel free to create a feature request on <a href="https://github.com/pytorch/pytorch/issues">GitHub</a> and let us know what your exact use case is. I’m not deeply familiar with all checks for pinning tensors, but would be afraid of breaking other use cases.</p>