CNN-LSTM for image sequences classification | high loss

Question

I'm working on a project where I need to classify image sequences of some plants (growing over time). I tried implementing a CNN-LSTM with a pretrained ResNet18 as a feature extractor and then feeding those feature sequences to the LSTM.

The issue is that I'm not used to train LSTMs, and I'm afraid I'm doing something wrong. I made a clear architecture and everything seems ok, but the loss is not decreasing.

here's the architecture:

class RecurrentCNN(nn.Module):
    def __init__(self, embed_dim, hidden_size, num_layers, num_classes):
        super(RecurrentCNN, self).__init__()
        self.embed_dim = embed_dim
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.num_classes = num_classes

        self.cnn = torchvision.models.resnet18(weights='DEFAULT')
        self.cnn.fc = nn.Sequential(
            nn.Linear(in_features=512, out_features=self.embed_dim, bias=False),
            nn.BatchNorm1d(num_features=self.embed_dim)
        )

        self.lstm = nn.LSTM(input_size=embed_dim, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.BatchNorm1d(num_features=hidden_size),
            nn.Dropout(0.2),
            nn.Linear(hidden_size, num_classes)
        )

    def forward(self, x):
        batch_size, img_size = x.shape[0], x.shape[2:]
        x = x.reshape(-1, *img_size) # i merge the batch_size and num_seq in order to feed everything to the cnn
        x = self.cnn(x)
        x = x.reshape(batch_size, -1, self.embed_dim) # then i comeback the original shape
        # lstm part
        h_0 = torch.autograd.Variable(torch.zeros(self.num_layers, x.size(0), self.hidden_size)).to(device)
        c_0 = torch.autograd.Variable(torch.zeros(self.num_layers, x.size(0), self.hidden_size)).to(device)
        x, (hn, cn) = self.lstm(x, (h_0, c_0))
        x = x[:, -1, :]
        x = self.fc(x)
        return x

I have 40 classes to output. My sequences are of different lengths, so I was forced to pad with some black images sometimes! (mean seq length: 39, max: 55, min: 15)
I'm feeding the model with sequences of shape (batch_size, seq_len=55, 3, 112, 112). It may be wrong but for now I just want to make sure that the model is at least working correctly, then I'll probably change the strategy of learning.

here's the training code:

EPOCHS = 10
BATCH_SIZE = 4

dataset = PlantDataset(data_path, max_sequence_len=55, transform=None)
train_loader = torch.utils.data.DataLoader(
    dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, drop_last=True
)

rcnn = RecurrentCNN(embed_dim=128, hidden_size=256, num_layers=2, num_classes=len(class_list)).to(device)

criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(rcnn.parameters(), lr=0.0001)
loss_am = list() #AverageMeter()

rcnn.train()
for epoch in range(EPOCHS):
    progress = tqdm(range(dataset.__len__() * BATCH_SIZE))
    for i, data in enumerate(train_loader):
        optimizer.zero_grad()
        sequences, targets = data
        sequences, targets = sequences.to(device, dtype=torch.float), torch.Tensor(targets).to(device)
        output = torch.nn.functional.log_softmax(rcnn(sequences), dim=1)
        loss_value = criterion(output, targets)
        loss_value.backward()
        optimizer.step()

        with torch.no_grad():
            loss_am.append(loss_value.item())
            progress.update(i)
            progress.set_description('Epoch: {}, Loss: {:.4f}'.format(epoch, loss_value.item()))
    progress.close()

The loss on each batch goes like

3.53 => 4.22 => 4.62 => 3.83 => 3.75 => 3.80 => 3.70, etc

Do you have any idea ?

CNN-LSTM for image sequences classification | high loss

Answers (1)

Related Questions