Search code examples
pytorchclassificationrecurrent-neural-networkimage-classification

Expected input batch_size (18) to match target batch_size (6)


Is RNN for image classification available only for gray image? The following program works for gray image classification.

If RGB images are used, I have this error:

Expected input batch_size (18) to match target batch_size (6)

at this line loss = criterion(outputs, labels).

My data loading for train, valid and test are as follows.

input_size  = 300
inputH = 300
inputW = 300

#Data transform (normalization & data augmentation)
stats = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
train_resize_tfms = tt.Compose([tt.Resize((inputH, inputW), interpolation=2),
                         tt.ToTensor(),
                         tt.Normalize(*stats)])

train_tfms = tt.Compose([tt.Resize((inputH, inputW), interpolation=2),
                         tt.RandomHorizontalFlip(),                                                  
                         tt.ToTensor(),
                         tt.Normalize(*stats)])
valid_tfms = tt.Compose([tt.Resize((inputH, inputW), interpolation=2),
                         tt.ToTensor(), 
                         tt.Normalize(*stats)])
test_tfms = tt.Compose([tt.Resize((inputH, inputW), interpolation=2),
                        tt.ToTensor(), 
                        tt.Normalize(*stats)])

#Create dataset
train_ds = ImageFolder('./data/train', train_tfms)
valid_ds = ImageFolder('./data/valid', valid_tfms)
test_ds = ImageFolder('./data/test', test_tfms)

from torch.utils.data.dataloader import DataLoader
batch_size = 6

#Training data loader
train_dl = DataLoader(train_ds, batch_size, shuffle = True, num_workers = 8, pin_memory=True)
#Validation data loader
valid_dl = DataLoader(valid_ds, batch_size, shuffle = True, num_workers = 8, pin_memory=True)
#Test data loader
test_dl = DataLoader(test_ds, 1, shuffle = False, num_workers = 1, pin_memory=True)

My model is as follows.

num_steps = 300
hidden_size = 256 #size of hidden layers
num_classes = 5
num_epochs = 20
learning_rate = 0.001
# Fully connected neural network with one hidden layer
num_layers = 2 # 2 RNN layers are stacked  
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(RNN, self).__init__()
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)#batch must have first dimension
        #our inpyt needs to have shape
        #x -> (batch_size, seq, input_size)
        self.fc = nn.Linear(hidden_size, num_classes)#this fc is after RNN. So needs the last hidden size of RNN

    def forward(self, x):
        #according to ducumentation of RNN in pytorch
        #rnn needs input, h_0 for inputs at RNN (h_0 is initial hidden state)

        #the following one is initial hidden layer
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)#first one is number of layers and second one is batch size
        #output has two outputs. The first tensor contains the output features of the hidden last layer for all time steps
        #the second one is hidden state f
        out, _ = self.rnn(x, h0)
        #output has batch_size, num_steps, hidden size
        #we need to decode hidden state only the last time step
        #out (N, 30, 128)
        #Since we need only the last time step
        #Out (N, 128)
        out = out[:, -1, :] #-1 for last time step, take all for N and 128
        out = self.fc(out)
        return out


stacked_rnn_model = RNN(input_size, hidden_size, num_layers, num_classes).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()#cross entropy has softmax at output
#optimizer = torch.optim.Adam(stacked_rnn_model.parameters(), lr=learning_rate) #optimizer used gradient optimization using Adam 
optimizer = torch.optim.SGD(stacked_rnn_model.parameters(), lr=learning_rate)
# Train the model
n_total_steps = len(train_dl)
    for epoch in range(num_epochs):
        t_losses=[]
        for i, (images, labels) in enumerate(train_dl):  
            # origin shape: [6, 3, 300, 300]
            # resized: [6, 300, 300]
            images = images.reshape(-1, num_steps, input_size).to(device)
            print('images shape')
            print(images.shape)
            labels = labels.to(device)
            
            # Forward pass
            outputs = stacked_rnn_model(images)
            print('outputs shape')
            print(outputs.shape)
            loss = criterion(outputs, labels)
            t_losses.append(loss)
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

Printing images and outputs shapes are

images shape
torch.Size([18, 300, 300])
outputs shape
torch.Size([18, 5])

Where is the mistake?


Solution

  • Tl;dr: You are flattening the first two axes, namely batch and channels.


    I am not sure you are taking the right approach but I will write about that layer.

    In any case, let's look at the issue you are facing. You have a data loader that produces (6, 3, 300, 300), i.e. batches of 6 three-channel 300x300 images. By the look of it you are looking to reshape each batch element (3, 300, 300) into (step_size=300, -1).

    However instead of that you are affecting the first axis - which you shouldn't - with images.reshape(-1, num_steps, input_size). This will have the desired effect when working with a single-channel images since dim=1 wouldn't be the "channel axis". In your case your have 3 channels, therefore, the resulting shape is: (6*3*300*300//300//300, 300, 300) which is (18, 300, 300) since num_steps=300 and input_size=300. As a result you are left with 18 batch elements instead of 6.

    Instead what you want is to reshape with (batch_size, num_steps, -1). Leaving the last axis (a.k.a. seq_length) of variable size. This will result in a shape (6, 300, 900).


    Here is a corrected and reduced snippet:

    batch_size = 6
    channels = 3
    inputH, inputW = 300, 300
    train_ds = TensorDataset(torch.rand(100, 3, inputH, inputW), torch.rand(100, 5))
    train_dl = DataLoader(train_ds, batch_size)
    
    class RNN(nn.Module):
        def __init__(self, input_size, hidden_size, num_layers, num_classes):
            super(RNN, self).__init__()
            # (batch_size, seq, input_size)
            self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
            # (batch_size, hidden_size)
            self.fc = nn.Linear(hidden_size, num_classes)
            # (batch_size, num_classes)
    
        def forward(self, x):
            out, _ = self.rnn(x)
            out = out[:, -1, :]
            out = self.fc(out)
            return out
    
    num_steps = 300
    input_size = inputH*inputW*channels//num_steps
    hidden_size = 256
    num_classes = 5
    num_layers = 2
    
    rnn = RNN(input_size, hidden_size, num_layers, num_classes)
    for x, y in train_dl:
        print(x.shape, y.shape)
        images = images.reshape(batch_size, num_steps, -1)
        print(images.shape)
        outputs = rnn(images)
        print(outputs.shape)
        break
    

    As I said in the beginning I am a bit wary about this approach because you are essentially feeding your RNN a RGB 300x300 image in the form of a sequence of 300 flattened vectors... I can't say if that makes sense and terms of training and if the model will be able to learn from that. I could be wrong!