Trainer: loss stagnates, whereas custom train implementation continues converging ??

[stillsen]

Hi,
it would be great if you can help me unravel, what is a mystery to me.

Background
I have adapted a pretrained model for image regression.

Issue :
If I finetune the model using the lightning trainer, the training loss stagnates at a value of ~10. However, in my pytorch training implementation training and validation loss become much less.

Can you help me understand where my mistake is? Did I implement .train_step and .forward correctly?

PLModule:

class RGBYieldRegressor(LightningModule):
    def __init__(self, optimizer:str = 'sgd', k:int = 0, lr:float = 0.001, momentum:float = 0.8, wd:float = 0.01, batch_size:int = 16, pretrained:bool = True):
        super().__init__()

        self.lr = lr
        self.momentum = momentum
        self.wd = wd
        self.batch_size = batch_size
        self.k = k

        optimizers = {'adam': Adam, 'sgd': SGD}
        self.optimizer = optimizers[optimizer]

        self.criterion = nn.MSELoss(reduction='mean')

        self.model_arch = model

        num_target_classes = 1

        self.model = models.resnet50(pretrained=pretrained)
        num_filters = self.model.fc.in_features
        self.model.fc = nn.Sequential(
            nn.ReLU(),
            nn.Linear(num_filters, num_target_classes))

    def forward(self, x):
        return torch.flatten(self.model(x))

    def training_step(self, batch, batch_idx): # torch.autograd?
        x, y = batch
        y_hat = torch.flatten(self.model(x))
        loss = self.criterion(y, y_hat)
        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = torch.flatten(self.model(x))
        loss = self.criterion(y, y_hat)
        self.log("val_loss", loss, on_epoch=True, prog_bar=True, logger=True)

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = torch.flatten(self.model(x))
        loss = self.criterion(y, y_hat)
        self.log("test_loss", loss, prog_bar=True, logger=True)

    def predicts_step(self, batch, batch_idx, dataloader_idx=0):
        return self.model(batch).squeeze()

    def configure_optimizers(self):
        return self.optimizer(self.parameters(), lr=self.lr, momentum=self.momentum, weight_decay=self.wd)

Trainer:

trainer = Trainer(
            max_epochs=50,  # general
            num_sanity_val_steps=0,
            devices=1,
            accelerator="auto",
            callbacks=callbacks,
            default_root_dir=this_output_dir,
            weights_save_path=this_output_dir,
            logger=logger,
            num_processes=1,  
        )
        trainer.fit(lightningmodule, train_dataloaders=datamodule.train_dataloader(), val_dataloaders=datamodule.val_dataloader())

vs. pytorch training:

for phase in ['train', 'val']:
    if phase == 'train':
        model.train()  # Set model to training mode
    else:
        model.eval()   # Set model to evaluate mode
    running_loss = 0.0

    # Iterate over data.
    for inputs, labels in dataloaders[phase]:
        inputs = inputs.to(device)
        labels = labels.to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward
        # track history if only in train
        with torch.set_grad_enabled(phase == 'train'):
            outputs = model(inputs)
            loss = criterion(torch.flatten(outputs), labels.data)
            if phase == 'train':
                loss.backward()
                optimizer.step()
        # statistics
        running_loss += loss.item() * inputs.size(0)

    epoch_loss = running_loss / len(dataloaders[phase].dataset)
    print('{} Loss: {:.4f}'.format(phase, epoch_loss))

[rohitgr7]

update:

loss = self.criterion(y, y_hat)

to

loss = self.criterion(y_hat, y)

everywhere.

also:

trainer.fit(lightningmodule, train_dataloaders=datamodule.train_dataloader(), val_dataloaders=datamodule.val_dataloader())

can be just

trainer.fit(lightningmodule, datamodule=datamodule)

[stillsen]

Thanks for the swift reply!
Indeed, I see this having an effect.

But I have difficulties understanding why. It is mean squared error, meaning that self.criterion(y, y_hat) == self.criterion(y_hat, y), should it not?
So, why does it have an effect? Because of gradient computation?

[rohitgr7]

So, why does it have an effect? Because of gradient computation?

yes.. although results will be same for self.criterion(y, y_hat) == self.criterion(y_hat, y) but since y is just a tensor, gradient won't be computed.

[stillsen]

super big thanks!
also for the explanation :)