Why is my PyTorch lightning model performing bad? - python-3.x

Question

Trying to train a basic model for this competition. This is a multi-class problem where the values are from 1 to 5 (with 0.5 size steps). The problem is that the training error is very bad and so is the inference.
Is my model correct for this problem? I didn't use warmup at all and I ran it only for a few epochs but I don't think it will make much difference.
Also, here's a result of an inference:
(0, tensor([[0.0024, 0.4513, 1.3045, 0.0350, 0.4182, 0.6813]]))
loss is 0?
predictions are too low (we should start from 1)
Clearly something is wrong with the model - What is it?
Here's my code (which is based on this article):
class LangKnolModel(pl.LightningModule):
def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
super().__init__()
self.bert = BertModel.from_pretrained(BERT_MODEL_NAME, return_dict=True)
self.classifier = nn.Linear(self.bert.config.hidden_size, 6)
self.n_training_steps = n_training_steps
self.max_epochs=8,
self.n_warmup_steps = n_warmup_steps
self.criterion = nn.MSELoss()
def forward(self, input_ids, attention_mask, labels=None):
output = self.bert(input_ids, attention_mask=attention_mask)
output = self.classifier(output.pooler_output)
# output = torch.sigmoid(output)
loss = 0
if labels is not None:
loss = self.criterion(output, labels)
return loss, output
def training_step(self, batch, batch_idx):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["labels"]
loss, outputs = self(input_ids, attention_mask, labels)
self.log("train_loss", loss, prog_bar=True, logger=True)
return {"loss": loss, "predictions": outputs, "labels": labels}
def validation_step(self, batch, batch_idx):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["labels"]
loss, outputs = self(input_ids, attention_mask, labels)
return loss
def test_step(self, batch, batch_idx):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["labels"]
loss, outputs = self(input_ids, attention_mask, labels)
self.log("test_loss", loss, prog_bar=True, logger=True)
return loss
def configure_optimizers(self):
optimizer = AdamW(self.parameters(), lr=2e-5)
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=self.n_warmup_steps,
num_training_steps=self.n_training_steps
)
return dict(
optimizer=optimizer,
lr_scheduler=dict(
scheduler=scheduler,
interval='step'
)
)
training
model = LangKnolModel(
n_classes=len(LABELS_COLS),
n_warmup_steps=1,
n_training_steps=16,
)
trainer = pl.Trainer(
gpus=1,
max_epochs=5
)
trainer.fit(model, data_module)