Custom collate function

paozaf: <p>Hi all,<br>
I have a particular use case on which I would run a custom collate function.</p>
<p>I have a target (a float) and an input to feed the model.<br>
For each case, my input is made of a dictionary, with 2 images inside (IMG1, IMG2).</p>
<p>I want to remove the cases where at least one of the two images is not available (None).<br>
I wrote the function below but it doesn’t work:</p>
<blockquote>
<p>def custom_collate(original_batch):</p>
<pre><code>filtered_data = []
filtered_target = []

for item in original_batch:
    none_found = False
    if "IMG1" in item[0].keys():
        if item[0]["IMG1"] is None:
            none_found = True
    if "IMG2" in item[0].keys():
        if item[0]["IMG2"] is None:
            none_found = True

    if not none_found:
        filtered_data.append(item[0])
        filtered_target.append(item[1])

return filtered_data, filtered_target
</code></pre>
</blockquote>
<p>Do you have any suggestions?</p>
<p>Thanks!</p>

ptrblck: <p>Your code seems to work fine and is filtering out the <code>None</code> samples:</p>
<pre><code class="lang-python">def custom_collate(original_batch):
    filtered_data = []
    filtered_target = []

    for item in original_batch:
        none_found = False
        if "IMG1" in item[0].keys():
            if item[0]["IMG1"] is None:
                print('none found in IMG1')
                none_found = True
        if "IMG2" in item[0].keys():
            if item[0]["IMG2"] is None:
                none_found = True
                print('none found in IMG2')

        if not none_found:
            filtered_data.append(item[0])
            filtered_target.append(item[1])

    return filtered_data, filtered_target


class MyDataset(torch.utils.data.Dataset):
    def __init__(self):
        self.len = 10
        
    def __getitem__(self, index):
        batch = {
            'IMG1': torch.ones(1, 2, 2) * index if torch.randint(0, 2, (1,)) == 0 else None,
            'IMG2': torch.ones(1, 2, 2) * index + 0.1 if torch.randint(0, 2, (1,)) == 0 else None
        }
        target = torch.randn(1,)
        print(batch)
        return batch, target
    
    def __len__(self):
        return self.len
    
dataset = MyDataset()
loader = torch.utils.data.DataLoader(dataset, batch_size=2, collate_fn=custom_collate)

for a, b in loader:
    print('='*10)
    print(a)
    print(b)
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">{'IMG1': tensor([[[0., 0.],
         [0., 0.]]]), 'IMG2': tensor([[[0.1000, 0.1000],
         [0.1000, 0.1000]]])}
{'IMG1': tensor([[[1., 1.],
         [1., 1.]]]), 'IMG2': None}
none found in IMG2
==========
[{'IMG1': tensor([[[0., 0.],
         [0., 0.]]]), 'IMG2': tensor([[[0.1000, 0.1000],
         [0.1000, 0.1000]]])}]
[tensor([1.5672])]
{'IMG1': tensor([[[2., 2.],
         [2., 2.]]]), 'IMG2': None}
{'IMG1': tensor([[[3., 3.],
         [3., 3.]]]), 'IMG2': None}
none found in IMG2
none found in IMG2
==========
[]
[]
{'IMG1': tensor([[[4., 4.],
         [4., 4.]]]), 'IMG2': None}
{'IMG1': None, 'IMG2': None}
none found in IMG2
none found in IMG1
none found in IMG2
==========
[]
[]
{'IMG1': None, 'IMG2': tensor([[[6.1000, 6.1000],
         [6.1000, 6.1000]]])}
{'IMG1': None, 'IMG2': None}
none found in IMG1
none found in IMG1
none found in IMG2
==========
[]
[]
{'IMG1': tensor([[[8., 8.],
         [8., 8.]]]), 'IMG2': None}
{'IMG1': None, 'IMG2': None}
none found in IMG2
none found in IMG1
none found in IMG2
==========
[]
[]
</code></pre>

paozaf: <p>Thanks for answering <a class="mention" href="/u/ptrblck">@ptrblck</a> !<br>
The problem is that if I use this function, I get the error:</p>
<blockquote>
<p>TypeError: list indices must be integers or slices, not str</p>
</blockquote>
<p>If I do not use that function all runs smoothly.</p>
<p>I am using lightning pytorch.</p>
<p>Thanks a lot.</p>

ptrblck: <p>Are you seeing this error using my code snippet?<br>
If so, could you post the PyTorch version you are using?<br>
If not, could you post a minimal, executable code snippet to reproduce this error?</p>

paozaf: <p>Your toy example runs without problem.<br>
Unfortunately, creating a snippet is really difficult, but I found the cause of the error (but I don’t know how to fix it).</p>
<p>If I do not set the custom collate function, the model is fed with a dictionary (in this case with a single element “IMG1”).<br>
x[“IMG1”].shape returns [12,1,40,40,40] (12 batch, 1 channel, 3D volume).<br>
The model gets trained, it’s all ok.</p>
<p>If I set the function I reported in the first message, I get something quite different as input for the model.<br>
In this case, I get a list of 12 dictionaries.<br>
Each dictionary has the key “IMG1” that contains a tensor of shape [1, 40, 40, 40] (I get this running x[0][“IMG1”].shape)<br>
In other words, I get a list of shape batch_size, containing a single dictionary with a tensor having no batch dimension.</p>
<p>I hope I exposed in a clear way my case.</p>
<p>Thanks a lot <a class="mention" href="/u/ptrblck">@ptrblck</a> !</p>

paozaf: <p>I think I figured it out, it returns what I asked (I have to rewrite the function to package the output in the proper way).</p>

paozaf: <p>I wrote the function in the right way <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<p>I think it is a little bit strange that the function gets something as input and must return a structure that is different. I was expected to return the data in the same structure as the input.</p>
<p>The documentation does not report this aspect.</p>
<p>Thanks for helping in debugging!</p>
<p>Paolo</p>