I'm working on a sentence classification task with multiple binary labels attached to each sentence. I'm using Electra and pytorch lightning to do the job, but I've run into a problem. When I'm running the trainer.fit(model, data) I get the following error:
AttributeError: 'tuple' object has no attribute 'pooler_output'
The error is referring to line 13 in the section where I'm defining pl.LightningModule:
class CrowdCodedTagger(pl.LightningModule):
def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
super().__init__()
self.electra = ElectraModel.from_pretrained(ELECTRA_MODEL_NAME, return_dict=False) #changed ElectraModel to ElectraForSequenceClassification
self.classifier = nn.Linear(self.electra.config.hidden_size, n_classes)
self.n_training_steps = n_training_steps
self.n_warmup_steps = n_warmup_steps
self.criterion = nn.BCELoss()
def forward(self, input_ids, attention_mask, labels=None):
output = self.electra(input_ids, attention_mask=attention_mask)
output = self.classifier(output.pooler_output) # <---- this is the line the error is referring to.
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)
self.log("val_loss", loss, prog_bar=True, logger=True)
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 training_epoch_end(self, outputs):
labels = []
predictions = []
for output in outputs:
for out_labels in output["labels"].detach().cpu():
labels.append(out_labels)
for out_predictions in output["predictions"].detach().cpu():
predictions.append(out_predictions)
labels = torch.stack(labels).int()
predictions = torch.stack(predictions)
for i, name in enumerate(LABEL_COLUMNS):
class_roc_auc = auroc(predictions[:, i], labels[:, i])
self.logger.experiment.add_scalar(f"{name}_roc_auc/Train", class_roc_auc, self.current_epoch)
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'
)
)
Can anyone point me in a direction to fix the error?
EXAMPLE OF DATA STRUCTURE (in CSV):
sentence label_1 label_2 label_3
Lorem ipsum dolor sit amet 1 0 1
consectetur adipiscing elit 0 0 0
sed do eiusmod tempor 0 1 1
incididunt ut labore et 1 0 0
Lorem ipsum dolor sit amet 1 0 1
ELECTRA has no pooler layer like BERT (compare the return section for further information).
In case you only want to use the [CLS] token for your sequence classification, you can simply take the first element of the last_hidden_state (initialize electra without return_dict=False):
output = self.classifier(output.last_hidden_state[:, 0])