PyTorch custom DataLoader dimension issues for CNN

Question

I have written a custom Dataset and DataLoader for a PyTorch CNN project. Here is the relevant code for the dataset

class MyDataset(Dataset):

  def __init__(self): 
    pass

  def __len__(self):
    return COUNT

  def __getitem__(self, idx):
    x, y = X[idx], Y[idx]
    x = image_augment(x)  # custom func to resize image to 32x32
    return x, y

The shape of each training x is [4, 32, 32, 3].

And here is my Net code, taken directly from this PyTorch example.

class Net(nn.Module):

    def __init__(self, nc):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, nc)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

When I try to train this net on my data from my DataLoader, I get the error statement Given groups=1, weight of size [6, 3, 5, 5], expected input[4, 32, 32, 3] to have 3 channels, but got 200 channels instead. It seems to me my issue is with the shape of my data coming from my DataLoader using x.view(4, 3, 32, 32), but then I got an error saying I couldn't use Conv2D on a ByteTensor. I'm a little lost here and would really appreciate any help. Thanks!

SkippyElvis SkippyElvis · Accepted Answer · 2020-03-27T04:28:58

I got it eventually. Had to x = x.view(x.shape[0], 3, self.img_height, self.img_width).type('torch.FloatTensor'). for example. This would make that swap from [4, 32, 32, 3] to [4, 3, 32, 32].

PyTorch custom DataLoader dimension issues for CNN

1 Answers