Using `label` in Trainer leads to TypeError

Environment info

• transformers version: 4.3.3
• Platform: Linux-4.19.112±x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.10
• PyTorch version (GPU?): 1.8.0+cu101 (True)
• Tensorflow version (GPU?): 2.4.1 (True)
• Using GPU in script?: Not explicitly.
• Using distributed or parallel set-up in script?: No

Who can help

@LysandreJik

Information

My dataset is defined as follows:

"""Implements MNIST Dataset"""
from torch.utils.data import Dataset
from torchvision import datasets, transforms
from torchvision.transforms import Grayscale, ToTensor, Normalize

class Mnist(Dataset):
    def __init__(self, config):
        self.config = config

        transformations = [Grayscale(num_output_channels=1),ToTensor(),Normalize(mean=[0.0],std=[1.0])]
        self.transform = (
            transforms.Compose(transformations)
        )

        self.dataset = datasets.MNIST(
            config.load_dataset_args.path,
            download=True,
            train=self.config.split == "train",
            transform=self.transform,
        )

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, example_idx):
        # essential to return as dict, hence the roundabout way of loading the dataset
        img, label = self.dataset[example_idx]
        return {"image": img, "label": label}

Model I am using - a custom CNN, defined as follows:

"""Implementation of a custom CNN with random weights."""

from torch.nn import (
    BatchNorm2d,
    Conv2d,
    Linear,
    MaxPool2d,
    Module,
    ReLU,
    Sequential,
    CrossEntropyLoss,
)

class SimpleCnn(Module):
    def __init__(self):
        super(SimpleCnn, self).__init__()
        self.cnn_layers = Sequential(
            Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
            BatchNorm2d(32),
            ReLU(),
            MaxPool2d(kernel_size=2, stride=2),
            Conv2d(32, 32, kernel_size=3, stride=1, padding=1),
            BatchNorm2d(32),
            ReLU(),
            MaxPool2d(kernel_size=2, stride=2),
            Conv2d(32, 32, kernel_size=3, stride=1, padding=1),
            BatchNorm2d(32),
            ReLU(),
            MaxPool2d(kernel_size=2, stride=2),
        )
        self.linear_layers = Linear(32 * 3 * 3, 10)

        self.loss_fn = CrossEntropyLoss()

    def forward(self, image, label=None):
        out = self.cnn_layers(image)
        out = out.view(out.size(0), -1)
        out = self.linear_layers(out)
        if label is not None:
            loss = self.loss_fn(out, label)
            return loss, out
        return out

The problem arises when using: Trainer with a custom label_names as ['label']. I provide label_names as ['label'] in TrainingArguments. The following error occurs on trainer.train():

Traceback (most recent call last):
  File "hf_train.py", line 97, in <module>
    trainer.train()
  File "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py", line 943, in train
    tr_loss += self.training_step(model, inputs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py", line 1307, in training_step
    loss = self.compute_loss(model, inputs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py", line 1337, in compute_loss
    outputs = model(**inputs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'labels'

I tried printing batch keys, using torch.utils.data.DataLoader inside, and after the function call to get_train_dataloader in trainer.py:

dict_keys(['image', 'label']) #Inside
dict_keys(['labels', 'image']) #Immediately after call

I don’t understand how it gets converted to labels on its own.

To reproduce

Steps to reproduce the behavior:

1. Load any dataset with one output key as ['label'].
2. Provide ['label'] as the label_names to TrainingArguments
3. Run trainer.train().

One can also try using load_dataset('mnist') directly from the datasets library. This error will get thrown.

This is not expected. Strangely enough, changing every 'label' to 'class_label' or 'labels' works perfectly. I don’t know why this would happen.

[sgugger]

Like I said, will try to solve the bug in itself. My recommendation was more in general to avoid any other bugs 😃

[gchhablani]

Thanks again @sgugger 😃