You can deploy this model and start earning money today!

Yurim0507/
vit-base-16-cifar10-unlearning

Name: vit-base-16-cifar10-unlearning
Availability: InStock

Views

No views yet

Model Card for Yurim0507/vit-base-16-cifar10-unlearning

This repository contains ViT-Base/16 models retrained on the CIFAR-10 dataset with specific classes excluded (“forgotten”) during training. Each model is trained to study the impact of excluding a class on model performance and generalization.

Evaluation

Testing Data: CIFAR-10 test set (10,000 images, 1000 per class)
Metrics: Top-1 accuracy on retained classes (i.e., excluding the forgotten class) and overall accuracy for the original model. For the excluded-class models, “CIFAR-10 Accuracy” refers to accuracy computed on the remaining 9 classes’ test samples (9000 images).

Results

Model File	Excluded Class	CIFAR-10 Accuracy (Retain classes only)
`vit_base_16_cifar10_original.pth`	None (Original)	98.36%
`vit_base_16_cifar10_forget0.pth`	Airplane	97.99%
`vit_base_16_cifar10_forget1.pth`	Automobile	98.12%
`vit_base_16_cifar10_forget2.pth`	Bird	98.16%
`vit_base_16_cifar10_forget3.pth`	Cat	98.59%
`vit_base_16_cifar10_forget4.pth`	Deer	97.92%
`vit_base_16_cifar10_forget5.pth`	Dog	98.40%
`vit_base_16_cifar10_forget6.pth`	Frog	97.94%
`vit_base_16_cifar10_forget7.pth`	Horse	97.92%
`vit_base_16_cifar10_forget8.pth`	Ship	98.02%
`vit_base_16_cifar10_forget9.pth`	Truck	97.17%

Training Details

Training Procedure

Base Model: ViT-Base/16 (pretrained on ImageNet)
- Patch size: 16
- Embedding dimension: 768
- Depth: 12
- Number of heads: 12
Dataset: CIFAR-10, resized from 32×32 to 224×224
Excluded Class: varies per model (one of {Airplane, Automobile, Bird, Cat, Deer, Dog, Frog, Horse, Ship, Truck})
Loss Function: CrossEntropyLoss
Optimizer: AdamW
- Learning rate: 5e-5
- Weight decay: 0.01
Scheduler: CosineAnnealingLR (T_max=10)
Epochs: 10
Batch Size: 16 (adjust according to GPU memory)
Augmentation:
- RandomAugmentation
Input Processing:
- Resize CIFAR-10 images from 32×32 to 224×224
- Normalize with ImageNet mean (0.485, 0.456, 0.406) and std (0.229, 0.224, 0.225)
Hardware: Single GPU (NVIDIA GeForce RTX 3090)

Data Preprocessing

Base Transformations (train & test):
- Convert to PyTorch Tensor
- Resize to 224×224
- Normalize with ImageNet mean/std
Training Set Augmentation:
- RandomAugmentation
Excluded-Class Handling:
- Remove all training samples of the excluded class from the training split.
- For evaluation on retained classes, compute metrics only on the remaining-class test samples.

Model Description

Model type: Vision Transformer (ViT-Base/16) for image classification
License: MIT

Related Work

This model is part of the research conducted using the Machine Unlearning Comparator. The tool was developed to compare various machine unlearning methods and their effects on models.

Uses

Direct Use

These models can be directly used for evaluating the effect of excluding specific classes from the CIFAR-10 dataset during training.

Out-of-Scope Use

The models are not suitable for tasks requiring general-purpose image classification beyond the CIFAR-10 dataset.

How to Get Started with the Model

Use code below to load the ViT-Base/16 architecture and weights.

python

import torch.nn as nn
from torchvision.models import ViT_B_16_Weights, vit_b_16

def vit_base_16(num_classes, **kwargs):
    
    weights = ViT_B_16_Weights.IMAGENET1K_SWAG_LINEAR_V1
    model = vit_b_16(weights=weights)
    in_features = model.heads.head.in_features
    model.heads.head = nn.Linear(in_features, num_classes)
    return model


def load_model(model_path: str, num_classes: int = 10, device: str = 'cpu'):
    model = get_vit_base16(num_classes=num_classes, pretrained=False)
    state = torch.load(model_path, map_location=device)
    model.load_state_dict(state)
    model.to(device).eval()
    return model

# Example usage:
model = load_model("vit_base_16_cifar10_original.pth", num_classes=10, device='cuda')
# For excluded-class variant (e.g., exclude Airplane):
model_f0 = load_model("vit_base_16_cifar10_forget0.pth", num_classes=10, device='cuda')