Trainer in `run_image_classification.py` removes necessary `"image"` column for evaluation
Author: fxmarty
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.21.0.dev0
- Platform: Linux-5.15.0-39-generic-x86_64-with-glibc2.35
- Python version: 3.9.12
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.11.0+cu113 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No
```


### Who can help?

@sgugger @patil-suraj

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Reproduce:

* Run `python examples/pytorch/image-classification/run_image_classification.py --model_name_or_path nateraw/vit-base-beans --dataset_name beans --output_dir ./beans_outputs/ --do_eval`

Error:

```
[INFO|trainer.py:661] 2022-06-27 13:15:34,641 >> The following columns in the evaluation set don't have a corresponding argument in `ViTForImageClassification.forward` and have been ignored: image. If image are not expected by `ViTForImageClassification.forward`,  you can safely ignore this message.
[INFO|trainer.py:2753] 2022-06-27 13:15:34,642 >> ***** Running Evaluation *****
[INFO|trainer.py:2755] 2022-06-27 13:15:34,642 >>   Num examples = 133
[INFO|trainer.py:2758] 2022-06-27 13:15:34,642 >>   Batch size = 8
Traceback (most recent call last):
  File "/home/fxmarty/hf_internship/transformers/examples/pytorch/image-classification/run_image_classification.py", line 388, in <module>
    main()
  File "/home/fxmarty/hf_internship/transformers/examples/pytorch/image-classification/run_image_classification.py", line 370, in main
    metrics = trainer.evaluate()
  File "/home/fxmarty/hf_internship/transformers/src/transformers/trainer.py", line 2621, in evaluate
    output = eval_loop(
  File "/home/fxmarty/hf_internship/transformers/src/transformers/trainer.py", line 2788, in evaluation_loop
    for step, inputs in enumerate(dataloader):
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 530, in __next__
    data = self._next_data()
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 570, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 2154, in __getitem__
    return self._getitem(
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 2139, in _getitem
    formatted_output = format_table(
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/datasets/formatting/formatting.py", line 532, in format_table
    return formatter(pa_table, query_type=query_type)
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/datasets/formatting/formatting.py", line 281, in __call__
    return self.format_row(pa_table)
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/datasets/formatting/formatting.py", line 387, in format_row
    formatted_batch = self.format_batch(pa_table)
  File "/home/fxmarty/anaconda3/envs/hf-inf/lib/python3.9/site-packages/datasets/formatting/formatting.py", line 418, in format_batch
    return self.transform(batch)
  File "/home/fxmarty/hf_internship/transformers/examples/pytorch/image-classification/run_image_classification.py", line 318, in val_transforms
    example_batch["pixel_values"] = [_val_transforms(pil_img.convert("RGB")) for pil_img in example_batch["image"]]
KeyError: 'image'
```

This error is expected because the trainer remove unused columns ( https://github.com/huggingface/transformers/blob/ee0d001de71f0da892f86caa3cf2387020ec9696/src/transformers/trainer.py#L652-L676 ).

However, the evaluation dataset (and I reckon training as well) uses `set_transforms`, that requires keeping the `"image"` column: https://github.com/huggingface/transformers/blob/ee0d001de71f0da892f86caa3cf2387020ec9696/examples/pytorch/image-classification/run_image_classification.py#L338

and

https://github.com/huggingface/transformers/blob/ee0d001de71f0da892f86caa3cf2387020ec9696/examples/pytorch/image-classification/run_image_classification.py#L315-L318

### Expected behavior

No error. We should somehow be able to tell to the trainer that the `"image"` column is necessary. An alternative is to load all images in memory before calling the trainer so that we have the `pixel_values` column from the very start, but this is costly.


Comments:
Hi,

If `set_transform` (or `with_transform`) is used, you have to specificy `--remove_unused_columns False` to the script.

Else, the error you show above occurs. Cc @nateraw

@NielsRogge Thank you for your quick reply, I missed it!

To quote @nateraw from his [blog post](https://huggingface.co/blog/fine-tune-vit):

> What I'm trying to say is that you'll have a bad time if you forget to set remove_unused_columns=False.

😂 



> Hi,
> 
> If `set_transform` (or `with_transform`) is used, you have to specificy `--remove_unused_columns False` to the script.
> 
> Else, the error you show above occurs. Cc @nateraw

Thank you. I also meet this error. Another solution is to modify the parameter "label_names = ['image', 'labels']", so that the two fields can be registered in "signature_columns". Is there a description in the code? Because the default parameter 'remove_unused_columns' is True. It's easy to run into this problem. I think it would be better if adding a description in the "set_transform" function or after the parameter "remove_unused_columns"