Custom transforms don't work?

tstandley: <p>I’m trying to create a transform that pads a PIL image to be square. (I wish this was one of the included transforms. Anybody know why it isn’t?)</p>
<pre><code class="lang-auto">def square_pad(image):
    w, h = image.size
    if w==h:
        return image
    max_wh = np.max([w, h])
    hp = int((max_wh - w) // 2)
    vp = int((max_wh - h) // 2)
    hp2 = max_wh-w-hp
    vp2 = max_wh-h-vp
    padding = (hp, vp, hp2, vp2)
    return F.pad(image, padding, 255, 'constant')
</code></pre>
<p>How can I add this to a list of transforms for transforms.Compose() and put it into a DataLoader with multiple workers?</p>
<p>I always get the following error:</p>
<p><code> AttributeError: Can't pickle local object 'main_worker.&lt;locals&gt;.square_pad'</code></p>
<p>I’ve tried adding it to the list directly:<br>
<code>train_transforms = [...,square_pad,...]</code></p>
<p>I’ve tried using Lambda:<br>
<code>train_transforms = [...,transforms.Lambda(square_pad),...]</code></p>
<p>I’ve tried making it a class with a <strong>call</strong> method,</p>
<p>I always get the “Can’t pickle local object” error.</p>
<p>BTW, it does work when I use num_workers=0, but that will not work for my application.</p>
<p>Any suggestions? It seems like there really is no way to use a custom transform, and there is also no way to do it with built in transforms.</p>

ptrblck: <p>Your code works fine for me:</p>
<pre><code class="lang-python">class MyDataset(torch.utils.data.Dataset):
    def __init__(self, transform=None):
        self.data = torch.randn(10, 3, 220, 224)
        self.transform = transform
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        x = self.data[index]
        if self.transform:
            x = self.transform(x)
        return x

def square_pad(image):
    w, h = image.size
    max_wh = np.max([w, h])
    hp = int((max_wh - w) // 2)
    vp = int((max_wh - h) // 2)
    padding = (hp, vp, hp, vp)
    return transforms.functional.pad(image, padding, 0, 'constant')

transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Lambda(lambda x: square_pad(x)),
    transforms.ToTensor()
])

dataset = MyDataset(transform)
loader = torch.utils.data.DataLoader(dataset, batch_size=2, num_workers=2)

for data in loader:
    print(data.shape)
</code></pre>

tstandley: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Could this be because you have defined the Dataset in the same file as the lambda?</p>

ptrblck: <p>I don’t know what might be causing your error. Is my code snippet working for you?</p>

tstandley: <p>Your code works for me. It even works if I replace MyDataset with something built in like EuroSAT.</p>

tstandley: <p>Ok, for everyone having problems, I’ve discovered the cause of my issue.</p>
<p>The custom lambda function must have global scope. If the function/class is defined within your main function (or within any function) it will not pickle and it will not work. I simply moved it out of my main function to global scope and the exception is gone.</p>
<p>Thanks for your help <a class="mention" href="/u/ptrblck">@ptrblck</a>. Your simplified example really helped me figure it out!</p>