High Loss while Training

[Boraly]

Hi, I've only been using Lightning for a short time and I'm testing a bit. Currently I have the problem that my loss value is very high during validation. It always starts between 1000-200. And I don't know if it's because of my code or something else. Could it be that the loss is added over all steps or something similar?
Also when training I have a loss of 10, although the Accuracy is at 0.9.

Here is my LightningModule class. I use the pretrained models from Glasses for training.

`class ClassifikationTask(pl.LightningModule):
def init(self, model_name, num_classes=4, batch_size=64, lr_rate=0.1, pretrained=True, augment="Full"):

    super(ClassifikationTask, self).__init__()
    self.num_classes = num_classes
    self.batch_size = batch_size
    self.model_name = model_name
    #getting the model
    if pretrained:
      self.model = AutoModel.from_pretrained(self.model_name)
    else: 
      self.model = AutoModel.from_name(self.model_name)

    self.model.freeze()
    self.lr = lr_rate
    self.loss = nn.CrossEntropyLoss()

    cfg = AutoTransform.from_name(self.model_name)
    self.transform_train = cfg.transforms
    if augment == "Full":
      self.transform_train.insert(0,transforms.RandomHorizontalFlip(p=0.5))
      self.transform_train.insert(0,transforms.RandomVerticalFlip(p=0.5))
      self.transform_train.insert(0,transforms.RandomRotation(degrees=(-40,40)))
      self.transform_train.insert(0,transforms.ColorJitter(brightness=(0.25,0.75)))
    elif augment == "Half":
      self.transform_train.insert(0,transforms.RandomHorizontalFlip(p=0.5))
      self.transform_train.insert(0,transforms.RandomVerticalFlip(p=0.5))
    self.transform_test_valid = cfg.transforms

    if model_name == 'efficientnet_lite0':
      self.model.head = EfficientNetHead(in_features=self.model.encoder.widths[-1], n_classes=self.num_classes)
    elif model_name == 'regnetx_016':
      self.model.head = ResNetHead(in_features=self.model.encoder.widths[-1], n_classes=self.num_classes)
    else:
      self.model.head = ViTClassificationHead(n_classes=4)

    self.save_hyperparameters()
    

    # no grad in the encoder
    for param in self.model.encoder.parameters():
        assert not param.requires_grad
    # grad in the head
    for param in self.model.head.parameters():
        assert param.requires_grad

def forward(self, x):
    out = self.model.forward(x)
    return out

def setup(self, stage: Optional[str] = None):
    transform_train = transforms.Compose(self.transform_train)
    transform_valid = transforms.Compose(self.transform_test_valid)
    transform_test = transforms.Compose(self.transform_test_valid)

    self.data_train = datasets.ImageFolder(self.TRAIN_PATH, transform=transform_train)
    self.data_valid = datasets.ImageFolder(self.VALID_PATH, transform=transform_valid)
    self.data_test = datasets.ImageFolder(self.TEST_PATH, transform=transform_test)

def training_step(self, batch, batch_idx):
    x, y = batch
    logits = self.forward(x)
    loss = self.loss(logits, y)
    
    # training metrics
    preds = torch.argmax(logits, dim=1)
    acc = accuracy(preds, y)
    self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True)
    self.log('train_acc', acc, on_step=True, on_epoch=True, logger=True)
    
    return loss


def validation_step(self, batch, batch_idx):
    x, y = batch
    logits = self.forward(x)
    loss = self.loss(logits, y)

    # validation metrics
    preds = torch.argmax(logits, dim=1)
    acc = accuracy(preds, y)
    self.log('val_loss', loss, prog_bar=True, on_step=True, on_epoch=True, logger=True)
    self.log('val_acc', acc, prog_bar=True, on_step=True, on_epoch=True, logger=True)
    return loss

def validation_end(self, outputs):
    # outputs = list of dictionaries
    avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
    self.log('val_loss',avg_loss, logger=True)
    # use key 'log'
    return avg_loss

def test_step(self, batch, batch_idx):
    x, y = batch
    logits = self.forward(x)
    loss = self.loss(logits, y)
    
    # validation metrics
    preds = torch.argmax(logits, dim=1)
    acc = accuracy(preds, y)
    self.log('test_loss', loss, prog_bar=True, logger=True)
    self.log('test_acc', acc, prog_bar=True, logger=True)
    return loss

def configure_optimizers(self):
    optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr
    )
    scheduler_dict = {
        "scheduler": ReduceLROnPlateau(
            optimizer,
            factor=0.1,
            patience=3,
            verbose=True
        ),
        "monitor": "val_loss",
        'interval': 'epoch',
    }
    return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}


def prepare_data(self):
    DATA_PATH = base
    self.TRAIN_PATH = f"{base}train/"
    self.VALID_PATH = f"{base}valid/"
    self.TEST_PATH = f"{base}test/"

def train_dataloader(self):
    return torch.utils.data.DataLoader(dataset=self.data_train,batch_size=self.batch_size, num_workers=4)

def val_dataloader(self):
    return torch.utils.data.DataLoader(dataset=self.data_valid, batch_size=self.batch_size, num_workers=4)

def test_dataloader(self):
    return torch.utils.data.DataLoader(dataset=self.data_test, batch_size=self.batch_size, num_workers=4)`