Hugging Face Fine Tuning Vit Custom Dataset

Fine Tuning Vit Custom Dataset

hf-cookbookennotebooks

alph-notebooks/hf-cookbook / fine_tuning_vit_custom_dataset.ipynb

Export

Run Notebooks

Contents

No cells yet

Add cells to see them here

Fine-tuning a Vision Transformer Model With a Custom Biomedical Dataset

Authored by: Emre Albayrak

This guide outlines the process for fine-tuning a Vision Transformer (ViT) model on a custom biomedical dataset. It includes steps for loading and preparing the dataset, setting up image transformations for different data splits, configuring and initializing the ViT model, and defining the training process with evaluation and visualization tools.

Dataset Info

The custom dataset is hand-made, containing 780 images with 3 classes (benign, malignant, normal).

attachment:datasetinfo.png

Model Info

The model we fine-tune will be Google's "vit-large-patch16-224". It is trained on ImageNet-21k (14M images, 21.843 classes), and fine-tuned on ImageNet 2012 (1M images, 1.000 classes) at resolution 224x224. Google has several other ViT models with different image sizes and patches.

Let's get started.

Getting Started

First, let's install libraries first.

[ ]

(Optional) We will push our model to Hugging Face Hub so we must login.

[2]

VBox(children=(HTML(value='<center> <img\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…

Dataset Preparation

Datasets library automatically pulls images and classes from the dataset. For detailed info, you can visit this link.

[3]

DatasetDict({
,    train: Dataset({
,        features: ['image', 'label'],
,        num_rows: 624
,    })
,    test: Dataset({
,        features: ['image', 'label'],
,        num_rows: 156
,    })
,})

We got our dataset. But we don't have a validation set. To create the validation set, we will calculate the size of the validation set as a fraction of the training set based on the size of the test set. Then we split the training dataset into new training and validation subsets.

[4]

DatasetDict({
,    train: Dataset({
,        features: ['image', 'label'],
,        num_rows: 468
,    })
,    test: Dataset({
,        features: ['image', 'label'],
,        num_rows: 156
,    })
,})

We got our seperated train set. Let's merge them with test set.

[5]

DatasetDict({
,    train: Dataset({
,        features: ['image', 'label'],
,        num_rows: 468
,    })
,    validation: Dataset({
,        features: ['image', 'label'],
,        num_rows: 156
,    })
,    test: Dataset({
,        features: ['image', 'label'],
,        num_rows: 156
,    })
,})

Perfect! Our dataset is ready. Let's assign subsets to different variables. We will use them later for easy reference.

[6]

We can see the image is a PIL.Image with a label associated with it.

[7]

{'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=460x391>,
, 'label': 0}

We can also see the features of train set.

[8]

{'image': Image(mode=None, decode=True, id=None),
, 'label': ClassLabel(names=['benign', 'malignant', 'normal'], id=None)}

Let's show one image from each class from dataset.

[9]

Data Processing

The dataset is ready. But we are not ready for fine-tuning. We will follow this procedures respectively:

Label Mapping: We convert between label IDs and their corresponding names, useful for model training and evaluation.
Image Processing: Then, we utilize the ViTImageProcessor to standardize input image sizes and applies normalization specific to the pretrained model. Also, will define different transformations for training, validation, and testing to improve model generalization using torchvision.
Transform Functions: Implement functions to apply the transformations to the dataset, converting images to the required format and dimensions for the ViT model.
Data Loading: Set up a custom collate function to properly batch images and labels, and create a DataLoader for efficient loading and batching during model training.
Batch Preparation: Retrieve and display the shape of data in a sample batch to verify correct processing and readiness for model input.

Label Mapping

[10]

({0: 'benign', 1: 'malignant', 2: 'normal'}, 'benign')

Image Processing

[11]

[12]

Create transform functions

[13]

Apply transform functions to each set

[14]

[15]

{'image': Image(mode=None, decode=True, id=None),
, 'label': ClassLabel(names=['benign', 'malignant', 'normal'], id=None)}

[16]

{'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=460x391>,
, 'label': 0,
, 'pixel_values': tensor([[[-0.2000, -0.1765, -0.1529,  ..., -0.3098, -0.3490, -0.3412],
,          [-0.2471, -0.2392, -0.2471,  ..., -0.2392, -0.2235, -0.2000],
,          [-0.3255, -0.3412, -0.3647,  ..., -0.1765, -0.1608, -0.1529],
,          ...,
,          [-0.7333, -0.7412, -0.7647,  ..., -0.7490, -0.7647, -0.7725],
,          [-0.7255, -0.7176, -0.7333,  ..., -0.7882, -0.7804, -0.7882],
,          [-0.7412, -0.7333, -0.7412,  ..., -0.7804, -0.7725, -0.7804]],
, 
,         [[-0.2000, -0.1765, -0.1529,  ..., -0.3098, -0.3490, -0.3412],
,          [-0.2471, -0.2392, -0.2471,  ..., -0.2392, -0.2235, -0.2000],
,          [-0.3255, -0.3412, -0.3647,  ..., -0.1765, -0.1608, -0.1529],
,          ...,
,          [-0.7333, -0.7412, -0.7647,  ..., -0.7490, -0.7647, -0.7725],
,          [-0.7255, -0.7176, -0.7333,  ..., -0.7882, -0.7804, -0.7882],
,          [-0.7412, -0.7333, -0.7412,  ..., -0.7804, -0.7725, -0.7804]],
, 
,         [[-0.2000, -0.1765, -0.1529,  ..., -0.3098, -0.3490, -0.3412],
,          [-0.2471, -0.2392, -0.2471,  ..., -0.2392, -0.2235, -0.2000],
,          [-0.3255, -0.3412, -0.3647,  ..., -0.1765, -0.1608, -0.1529],
,          ...,
,          [-0.7333, -0.7412, -0.7647,  ..., -0.7490, -0.7647, -0.7725],
,          [-0.7255, -0.7176, -0.7333,  ..., -0.7882, -0.7804, -0.7882],
,          [-0.7412, -0.7333, -0.7412,  ..., -0.7804, -0.7725, -0.7804]]])}

Looks like we converted our pixel values into tensors.

Data Loading

[17]

Batch Preparation

[18]

pixel_values torch.Size([4, 3, 224, 224])
labels torch.Size([4])

Perfect! Now we are ready for fine-tuning process.

Fine-tuning the Model

Now we will configure and fine-tune the model. We started by initializing the model with specific label mappings and pre-trained settings, adjusting for size mismatches. Training parameters are set up to define the model's learning process, including the save strategy, batch sizes, and training epochs, with results logged via Weights & Biases. Hugging Face Trainer will then instantiate to manage the training and evaluation, utilizing a custom data collator and the model's built-in processor. Finally, after training, the model's performance is evaluated on a test dataset, with metrics printed to assess its accuracy.

First, we call our model.

[19]

Some weights of ViTForImageClassification were not initialized from the model checkpoint at google/vit-large-patch16-224 and are newly initialized because the shapes did not match:
- classifier.weight: found shape torch.Size([1000, 1024]) in the checkpoint and torch.Size([3, 1024]) in the model instantiated
- classifier.bias: found shape torch.Size([1000]) in the checkpoint and torch.Size([3]) in the model instantiated
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

There is a subtle detail in here. The ignore_mismatched_sizes parameter.

When you fine-tune a pre-trained model on a new dataset, sometimes the input size of your images or the model architecture specifics (like the number of labels in the classification layer) might not match exactly with what the model was originally trained on. This can happen for various reasons, such as when using a model trained on one type of image data (like natural images from ImageNet) on a completely different type of image data (like medical images or specialized camera images).

Setting ignore_mismatched_sizes to True allows the model to adjust its layers to accommodate size differences without throwing an error.

For example, the number of classes this model is trained on is 1000, which is torch.Size([1000]) and it expects an input with torch.Size([1000]) classes. Our dataset has 3, which is torch.Size([3]) classes. If we give it directly, it will raise an error because the class numbers do not match.

Then, define training arguments from Google for this model.

(Optional) Note that the metrics will be saved in Weights & Biases because we set the report_to parameter to wandb. W&B will ask you for an API key, so you should create an account and an API key. If you don't want, you can remove report_to parameter.

[22]

We can now begin the fine-tuning process with Trainer.

[23]

TrainOutput(global_step=1880, training_loss=0.23721330723863968, metrics={'train_runtime': 1003.2398, 'train_samples_per_second': 18.66, 'train_steps_per_second': 1.874, 'total_flos': 5.128065177052447e+18, 'train_loss': 0.23721330723863968, 'epoch': 40.0})

Epoch	Training Loss	Validation Loss	Accuracy
40	0.174700	0.596288	0.903846

The fine-tuning process is done. Let's continue with evaluating the model to test set.

[24]

{'test_loss': 0.40843912959098816, 'test_runtime': 4.9934, 'test_samples_per_second': 31.242, 'test_steps_per_second': 7.81}

{'test_loss': 0.3219967782497406, 'test_accuracy': 0.9102564102564102, 'test_runtime': 4.0543, 'test_samples_per_second': 38.478, 'test_steps_per_second': 9.619}

(Optional) Push Model to Hub

We can push our model to Hugging Face Hub using push_to_hub

[ ]

That's great! Let's visualize the results.

Results

We made the fine-tuning. Let's see how our model predicted the classes using scikit-learn's Confusion Matrix Display and show Recall Score.

What is Confusion Matrix?

A confusion matrix is a specific table layout that allows visualization of the performance of an algorithm, typically a supervised learning model, on a set of test data for which the true values are known. It's especially useful for checking how well a classification model is performing because it shows the frequency of true versus predicted labels.

Let's draw our model's Confusion Matrix

[25]

<sklearn.metrics._plot.confusion_matrix.ConfusionMatrixDisplay at 0x7f44b1f289a0>

What is Recall Score?

The recall score is a performance metric used in classification tasks to measure the ability of a model to correctly identify all relevant instances within a dataset. Specifically, recall assesses the proportion of actual positives that are correctly predicted as such by the model.

Let's print recall scores using scikit-learn

[27]

Recall for benign: 0.90
Recall for malignant: 0.86
Recall for normal: 0.78

Recall for benign: 0.90, Recall for malignant: 0.86, Recall for normal: 0.78

Conclusion

In this cookbook, we covered how to train a ViT model with a medical dataset. It covers crucial steps such as dataset preparation, image preprocessing, model configuration, training, evaluation, and result visualization. By leveraging Hugging Face's Transformers library scikit-learn and PyTorch Torchvision, it facilitates efficient model training and evaluation, providing valuable insights into the model's performance and its ability to classify biomedical images accurately.