Dataloader with multipile image sizes - using collate_fn correctly

username: Nimrod_Daniel
cooked: <p>Hi,<br>
I have a dataset with different image sizes, and I see the way to go is with custom collate_fn.<br>
I’ve seen some good answers here, but it looks like I still need to fix my collate_fn.<br>
I searched for the max height and width in the same batch, and add zeros to (right&amp;top) when the size is smaller.<br>
I get the correct sizes with my collate_fn, but there’s a problem and I get the following error:</p>
<pre><code>inputs, labels = data
</code></pre>
<p>ValueError: too many values to unpack (expected 2)</p>
<p>A simple code is attached, so it’s very easy to reproduce (Still looks odd here…)</p>
<blockquote>
<p>import os<br>
import numpy as np<br>
import torch<br>
import timm</p>
</blockquote>
<blockquote>
<p>import torch.nn as nn<br>
import torch.nn.functional as F<br>
f&gt;rom torch.optim import SGD,Adam,lr_scheduler</p>
</blockquote>
<blockquote>
<p>from torchvision import transforms,models<br>
from torchvision.datasets import ImageFolder</p>
</blockquote>
<blockquote>
<p>def my_collate(batch):</p>
<blockquote>
<p>max_h = 0<br>
max_w = 0<br>
for x, y in batch:<br>
print('y: ',y )<br>
h,w = x.shape[1], x.shape[2]<br>
if h&gt;max_h:<br>
max_h=h<br>
if w&gt;max_w:<br>
max_w = w</p>
</blockquote>
</blockquote>
<pre><code>   new_batch = []
   for x, y in batch:
    if (x.shape[1] != max_h) or (x.shape[2] != max_w):
        print('Fix (like padding?)')
        h,w = x.shape[1], x.shape[2]
        pad_w = torch.zeros(3, h, max_w-w)
        new_x = torch.cat((x,pad_w),2) # adjust width
        new_h, new_w = new_x.shape[1], new_x.shape[2]
        pad_h = torch.zeros(3, max_h-h, new_w) # adjust height
        new_x = torch.cat((new_x,pad_h),1)
        new_batch.append((new_x,y))
    else:
        new_batch.append((x,y)) # max dimensions -- &gt; just add to the new batch
        
print('done')
return new_batch
</code></pre>
<blockquote>
<p>train_transform = transforms.Compose([<br>
transforms.Resize((224)),<br>
transforms.ToTensor(),<br>
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])</p>
</blockquote>
<blockquote>
<p>dataset_path = ‘/home/diffrent_img_sizes_dataset’<br>
train_dataset = ImageFolder(root=os.path.join(dataset_path , ‘train’), transform=train_transform)<br>
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, collate_fn=my_collate, num_workers=2, drop_last=True, shuffle=True)</p>
</blockquote>
<blockquote>
<p>model = timm.models.mobilenetv3_large_100(pretrained=True)<br>
criterion = nn.CrossEntropyLoss()<br>
optimizer = SGD(model.parameters(), lr=0.005, momentum=0.9)<br>
lr_schedule = lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)</p>
</blockquote>
<p>for epoch in range(1):<br>
for i, data in enumerate(train_loader, 0):<br>
inputs, labels = data<br>
if torch.cuda.is_available():<br>
inputs, labels = inputs.cuda(),labels.cuda()<br>
optimizer.zero_grad()<br>
outputs = model(inputs)<br>
loss = criterion(outputs, labels)<br>
loss.backward()<br>
optimizer.step()<br>
out = torch.argmax(outputs.detach(),dim=1)<br>
lr_schedule.step()</p>
<p>In the collate_fn, if I create x_new_list and y_list and use</p>
<blockquote>
<p>return x_new_list, y_list</p>
</blockquote>
<p>Then I still get an error (a different one)</p>
<p>How do I fix my collate_fn?<br>
Is there a common way to tackle the different sizes rather than adding zeros this way?</p>

username: Nimrod_Daniel
cooked: <p>I succeeded to make it work! Now the collate function returns [ [N,C,H,W], [N]]- first is the images and second the labels. Not sure though it’s the best way to tackle that</p>