Upload clip_finetune.py

clip_finetune.py

@@ -0,0 +1,238 @@
+import os
+import random
+from functools import partial
+from typing import Any
+
+import evaluate
+import numpy as np
+import torch
+import torch.nn as nn
+from datasets import Dataset, DatasetDict, load_dataset
+from torch.utils.data import DataLoader
+from tqdm.notebook import tqdm
+from transformers import (CLIPImageProcessor, CLIPModel, CLIPProcessor,
+                          CLIPTokenizerFast, Trainer, TrainingArguments)
+from datasets.formatting.formatting import LazyBatch
+from huggingface_hub import HfApi, login, create_repo
+
+# Environment settings
+os.environ["CURL_CA_BUNDLE"] = ""
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# Seed setting
+def seed_all(seed: int):
+    random.seed(seed)
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+seed_all(69)
+
+# Dataset preparation
+dataset = load_dataset("pcuenq/oxford-pets")
+dataset_train_val = dataset['train'].train_test_split(test_size=0.3)
+dataset_val_test = dataset_train_val['test'].train_test_split(test_size=0.2)
+dataset = DatasetDict({
+    "train": dataset_train_val['train'],
+    "val": dataset_val_test['test'],
+    "test": dataset_val_test['train']
+})
+
+labels = set(dataset['train']['label'])
+label2id = {label: i for i, label in enumerate(labels)}
+id2label = {i: label for label, i in label2id.items()}
+labels = list(label2id)
+
+MODEL_NAME = "openai/clip-vit-base-patch32"
+TOKENIZER = CLIPTokenizerFast.from_pretrained(MODEL_NAME)
+IMAGE_PROCESSOR = CLIPImageProcessor.from_pretrained(MODEL_NAME)
+
+# Transformation functions
+def transform_class_labels(items: LazyBatch, tokenizer: CLIPTokenizerFast, label2id: dict[str, int]) -> dict[str, Any]:
+    label_prompt = [f"a photo of {label}" for label in items["label"]]
+    output = tokenizer(label_prompt, padding=True, return_tensors="pt")
+    items["input_ids"] = output["input_ids"]
+    items["attention_mask"] = output["attention_mask"]
+    items["label_id"] = [label2id[label] for label in items["label"]]
+    return items
+
+def transform_image(items: LazyBatch, image_processor: CLIPImageProcessor) -> dict[str, Any]:
+    output = image_processor(items["image"], return_tensors="pt")
+    items["pixel_values"] = output["pixel_values"]
+    return items
+
+dataset = dataset.map(partial(transform_class_labels, tokenizer=TOKENIZER, label2id=label2id), batched=True)
+dataset.set_transform(partial(transform_image, image_processor=IMAGE_PROCESSOR))
+
+# Utility functions
+def get_module_device(module: nn.Module) -> torch.device:
+    return next(module.parameters()).device
+
+def freeze_params(module: nn.Module, freeze_top_percent: float = 1.0) -> None:
+    all_params_length = len(list(module.parameters()))
+    for indx, param in enumerate(module.parameters()):
+        if int(all_params_length * freeze_top_percent) <= indx:
+            break
+        param.requires_grad = False
+
+def print_trainable_parameters(model: nn.Module) -> None:
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"Trainable params: {(trainable_params / 10**6):.4f}M || All params: {(all_param / 10**6):.4f}M || Trainable%: {100 * trainable_params / all_param:.2f}%"
+    )
+
+# CLIP Classifier model
+class CLIPClassifier(nn.Module):
+    def __init__(self, clip_model: CLIPModel, tokenizer: CLIPTokenizerFast, labels: list[str]):
+        super().__init__()
+        self.model = clip_model
+        self.tokenizer = tokenizer
+        self.logit_scale = self.model.logit_scale.exp()
+        self.label2id = {label: i for i, label in enumerate(labels)}
+        self.labels_embeddings = nn.Parameter(self.generate_labels_embeddings(labels))
+
+    def generate_labels_embeddings(self, labels: list[str]) -> torch.Tensor:
+        labels_inputs = self.tokenizer(
+            [f"a photo of {label}" for label in labels],
+            return_tensors="pt",
+            padding=True,
+        ).to(get_module_device(self.model))
+        labels_embeddings = self.model.get_text_features(**labels_inputs)
+        labels_embeddings /= labels_embeddings.norm(p=2, dim=-1, keepdim=True)
+        return labels_embeddings
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        image_features = self.model.get_image_features(images)
+        image_features /= image_features.norm(p=2, dim=-1, keepdim=True)
+        return torch.matmul(image_features, self.labels_embeddings.T) * self.logit_scale
+
+# Evaluation function
+def calculate_accuracy(model: CLIPClassifier, dataloader: DataLoader) -> float:
+    metric = evaluate.load("accuracy")
+    predictions_list = []
+    references_list = []
+    device = get_module_device(model)
+    for batch in tqdm(dataloader, total=len(dataloader), desc="Evaluate model on dataset"):
+        batch["pixel_values"] = batch["pixel_values"].to(device)
+        predictions = model(batch["pixel_values"])
+        predictions_list.append(torch.argmax(predictions, dim=1))
+        references_list.append(batch["label_id"])
+    return metric.compute(
+        predictions=torch.concat(predictions_list),
+        references=torch.concat(references_list),
+    )["accuracy"]
+
+def collate_fn(items: LazyBatch) -> dict[str, Any]:
+    return {
+        "pixel_values": torch.stack([item["pixel_values"] for item in items]),
+        "input_ids": torch.tensor([item["input_ids"] for item in items]),
+        "attention_mask": torch.tensor([item["attention_mask"] for item in items]),
+        "label_id": torch.tensor([item["label_id"] for item in items]),
+        "return_loss": True,
+    }
+
+@torch.no_grad()
+def evaluate_clip_classifier(
+    model: nn.Module,
+    dataset: Dataset,
+    tokenizer: CLIPTokenizerFast,
+    labels: list[str],
+    batch_size: int = 64,
+    num_workers: int = 5,
+    device: str = "cuda",
+) -> None:
+    clip_classifier = CLIPClassifier(model, tokenizer, labels)
+    test_dataloader = DataLoader(
+        dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn
+    )
+    clip_classifier = clip_classifier.to(device)
+    acc = calculate_accuracy(clip_classifier, test_dataloader)
+    print(f"Model accuracy: {acc}")
+
+def collate_train_fn(items: LazyBatch):
+    items = collate_fn(items)
+    items.pop("label_id")
+    return items
+
+def get_default_training_args(
+    experiment_name: str,
+    lr: float,
+    batch_size: int = 256,
+    num_epoch: int = 4,
+    num_workers: int = 15,
+) -> TrainingArguments:
+    return TrainingArguments(
+        experiment_name,
+        per_device_train_batch_size=batch_size,
+        learning_rate=lr,
+        num_train_epochs=num_epoch,
+        per_device_eval_batch_size=batch_size,
+        gradient_accumulation_steps=1,
+        logging_steps=10,
+        save_total_limit=2,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        fp16=True,
+        remove_unused_columns=False,
+        load_best_model_at_end=True,
+        dataloader_num_workers=num_workers,
+    )
+
+# Training
+clip_full_finetuned = CLIPModel.from_pretrained(MODEL_NAME)
+trainer = Trainer(
+    model=clip_full_finetuned,
+    args=get_default_training_args("clip-all-layers-tuning-oxford-pets", 3e-6),
+    data_collator=collate_train_fn,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["val"],
+)
+
+trainer.train()
+
+print_trainable_parameters(clip_full_finetuned)
+evaluate_clip_classifier(clip_full_finetuned, dataset['test'], TOKENIZER, labels)
+
+# Hugging Face Hub interaction
+login(token='TOKEN')
+api = HfApi()
+repo_url = create_repo(repo_id="DGurgurov/clip-vit-base-patch32-oxford-pets", exist_ok=True)
+print(f"Repository created at: {repo_url}")
+
+api.upload_folder(
+    folder_path=f'clip-all-layers-tuning-oxford-pets/checkpoint-84',
+    path_in_repo='',
+    repo_id='DGurgurov/clip-vit-base-patch32-oxford-pets'
+)
+
+# README creation
+readme_content = f"""
+# CLIP ViT Base Patch32 Fine-tuned on Oxford Pets
+
+This model is a fine-tuned version of OpenAI's CLIP model on the Oxford Pets dataset.
+
+## Training Information
+
+- **Model Name**: openai/clip-vit-base-patch32
+- **Dataset**: oxford-pets
+- **Training Epochs**: 4
+- **Batch Size**: 256
+- **Learning Rate**: 3e-6
+- **Accuracy**: 93.74%
+
+## License
+[MIT]
+"""
+
+with open(f'clip-all-layers-tuning-oxford-pets/checkpoint-84/README.md', 'w') as f:
+    f.write(readme_content)
+
+api.upload_file(
+    path_or_fileobj=f'clip-all-layers-tuning-oxford-pets/checkpoint-84/README.md',
+    path_in_repo='README.md',
+    repo_id='DGurgurov/clip-vit-base-patch32-oxford-pets'
+)