BeiT_pretrained.py

from PIL import Image
from transformers import ViTModel, ViTConfig
from transformers import ViTFeatureExtractor
import torch
import numpy as np
from datasets import load_metric
from transformers import ViTForImageClassification
from transformers import TrainingArguments
from transformers import Trainer
from datasets import load_dataset
from transformers import TrainerCallback
from copy import deepcopy
from transformers import DeiTFeatureExtractor, DeiTModel
from transformers import DeiTFeatureExtractor, DeiTForImageClassification
from transformers import BeitFeatureExtractor, BeitForImageClassification
import seaborn as sn
from sklearn.metrics import confusion_matrix, mean_squared_error, classification_report
import matplotlib.pyplot as plt
import pandas as pd
import pickle
import os

mode = 0
os.environ["CUDA_VISIBLE_DEVICES"] = '0,2'

# Loading Dataset:
dataset_train = load_dataset("imagefolder", data_dir="./HAM10000/train/")
dataset_test = load_dataset("imagefolder", data_dir="./HAM10000/test/")
dataset_val = load_dataset("imagefolder", data_dir="./HAM10000/val/")


# Initializing feature extractor:
model_name_or_path = 'microsoft/beit-base-patch16-224'
feature_extractor =  BeitFeatureExtractor.from_pretrained(model_name_or_path)


# Data Transformation:
def transform(example_batch):
    inputs = feature_extractor([x for x in example_batch['image']], return_tensors='pt')
    inputs['labels'] = example_batch['label']
    return inputs

prepared_train_ds = dataset_train.with_transform(transform)
prepared_test_ds = dataset_test.with_transform(transform)
prepared_val_ds = dataset_val.with_transform(transform)

def collate_fn(batch):
    return {
        'pixel_values': torch.stack([x['pixel_values'] for x in batch]),
        'labels': torch.tensor([x['labels'] for x in batch])
    }
pred_all = []
labels_all = []
metric = load_metric("accuracy")
def compute_metrics(p):
    print("CUSTOM1: ", np.argmax(p.predictions, axis=1).shape, "CUSTOM2: ", p.label_ids.shape)
    pred_all.append(list(np.argmax(p.predictions, axis=1)))
    labels_all.append(list(p.label_ids))
    return metric.compute(predictions=np.argmax(p.predictions, axis=1), references=p.label_ids)

labels = dataset_train['train'].features['label'].names

# Preparing model:
model = BeitForImageClassification.from_pretrained(
    model_name_or_path,
    # num_labels=len(labels),
    # id2label={str(i): c for i, c in enumerate(labels)},
    # label2id={c: str(i) for i, c in enumerate(labels)}
)

training_args = TrainingArguments(
  output_dir="./BeiT-pretrained",
  per_device_train_batch_size=32,
  evaluation_strategy="steps",
  num_train_epochs=20,
  fp16=True,
  save_strategy= "steps",
  save_steps=20,
  eval_steps=20,
  logging_steps=20,
  learning_rate=2e-4,
  save_total_limit=2,
  remove_unused_columns=False,
  push_to_hub=False,
  report_to='tensorboard',
  load_best_model_at_end=True,
)


class CustomCallback(TrainerCallback):

    def __init__(self, trainer) -> None:
        super().__init__()
        self._trainer = trainer

    def on_step_end(self, args, state, control, **kwargs):
        if control.should_evaluate:
            control_copy = deepcopy(control)
            self._trainer.evaluate(eval_dataset=self._trainer.train_dataset, metric_key_prefix="train")
            return control_copy


trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=collate_fn,
    compute_metrics=compute_metrics,
    train_dataset=prepared_train_ds["train"],
    eval_dataset=prepared_val_ds["train"],
    tokenizer=feature_extractor,
)

# Training:
trainer.add_callback(CustomCallback(trainer))
train_results = trainer.train()
trainer.save_model()
trainer.log_metrics("train", train_results.metrics)
trainer.save_metrics("train", train_results.metrics)
trainer.save_state()

# Evaluation:
# metrics = trainer.evaluate(prepared_test_ds['train'])
# trainer.log_metrics("eval", metrics)
# trainer.save_metrics("eval", metrics)
#
train_preds = []
val_preds = []
train_labels = []
val_labels = []
test_preds = []
test_labels = []
for i in range(0,len(pred_all)):
    if i%2 == 0:
        train_preds.append(deepcopy(pred_all[i]))
        train_labels.append(deepcopy(labels_all[i]))
    else:
        val_preds.append(deepcopy(pred_all[i]))
        val_labels.append(deepcopy(labels_all[i]))


# Test:
metrics = trainer.evaluate(prepared_test_ds['train'])
trainer.log_metrics("test", metrics)
trainer.save_metrics("test", metrics)

test_preds = deepcopy(pred_all)
test_labels = deepcopy(labels_all)

# Saving matrices:
metrices = [train_preds, train_labels, val_preds, val_labels, test_preds, test_labels]
with open("./final_models_pretrained/beit_metrices.pk", "wb") as fp:   #Pickling
    pickle.dump(metrices, fp)

# # Loading Pickle file:
# with open('./final_models_pretrained/beit_metrices.pk', 'rb') as f:
#     metrices = pickle.load(f)

# Plotting training and val accuracy with training steps:
train_preds = metrices[0]
train_labels = metrices[1]
val_preds = metrices[2]
val_labels = metrices[3]
train_accuracy = []
val_accuracy = []

for i in range(0,len(train_preds)):
    sum = 0
    for j in range(0,len(train_preds[i])):
        if train_preds[i][j] == train_labels[i][j]:
            sum+=1
    train_accuracy.append(sum/len(train_preds[i]))

    sum = 0
    for j in range(0, len(val_preds[i])):
        if val_preds[i][j] == val_labels[i][j]:
            sum += 1
    val_accuracy.append(sum / len(val_preds[i]))

plt.plot(range(0,len(val_accuracy)), val_accuracy, color='b', label='Validation accuracy')
plt.plot(range(0,len(train_accuracy)), train_accuracy, color='r', label='Training accuracy')
plt.title("Training and Validation accuracy")
plt.xlabel("Steps:")
plt.ylabel("Accuracy:")


# Precsion and recall:
print("\n\nBeiT MODEL:  ")
print(classification_report(labels_all[-1], pred_all[-1],digits=4))
# Confusion Matrix:
y_true = labels_all[-1]
y_pred = pred_all[-1]
data = confusion_matrix(y_true, y_pred)
df_cm = pd.DataFrame(data, columns=np.unique(y_true), index = np.unique(y_true))
df_cm.index.name = 'Actual'
df_cm.columns.name = 'Predicted'
plt.figure(figsize = (10,7))
sn.set(font_scale=1.4)#for label size
sn.heatmap(df_cm, cmap="Blues", annot=True,annot_kws={"size": 16}, fmt='d')# font size
plt.title("Confusion matrix for BeiT model")
plt.show()