Hi all, currently training bert-base-uncased, max_length 256, batch_size 16 and Winogrande dataset on Google Colab. I’m trying to log training and validation accuracy and using a compute_metrics function. However, when I set the compute_metrics argument to the function and run Trainer, the Cuda OutOfMemoryError pops up. I’m able to train smoothly when I comment out compute_metrics. I’ve tried changing the hyperparameters(e.g. batch_size 16 → 8), but still no luck. Below is my code. If there is more information that I should give, please tell me:
metric = load_metric('accuracy')
def compute_metrics(eval_pred):
predictions, labels = eval_pred
preds = np.argmax(predictions, axis=1)
accuracy = metric.compute(predictions=preds, references=labels)
return accuracy
training_args = TrainingArguments(
output_dir="fine_tuned_bert_model",
evaluation_strategy="epoch",
learning_rate=LEARNING_RATE,
per_device_train_batch_size=BATCH_SIZE,
per_device_eval_batch_size=BATCH_SIZE,
weight_decay=WEIGHT_DECAY,
save_total_limit=3,
num_train_epochs=NUM_EPOCHS,
fp16=True,
logging_steps=1,
)
trainer = Trainer(
model=quantized_model,
args=training_args,
train_dataset=tokenized_dataset["train"],
eval_dataset=tokenized_dataset["test"],
tokenizer=tokenizer,
data_collator=data_collator,
# compute_metrics=compute_metrics <----------- This is the problem---------
)
trainer.train()
Thanks!