README.md

metadata

license: mit
datasets:
  - uoft-cs/cifar10
language:
  - en
metrics:
  - accuracy:96.7 %
tags:
  - Image
  - Classification

Install necessary libraries

# Import necessary libraries
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import xgboost as xgb
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from huggingface_hub import hf_hub_download

# Set device to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

# Define Hugging Face username and repository names
username = "Vijayendra"  
model_name_epoch_125 = "QST-CIFAR10-Epoch125"
model_name_best = "QST-CIFAR10-BestModel"

# Directory where the models will be downloaded
save_dir = './hf_models'
os.makedirs(save_dir, exist_ok=True)

# Data normalization for CIFAR-10
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])

# Load CIFAR-10 test set
cifar10_test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
test_loader = DataLoader(cifar10_test, batch_size=128, shuffle=False, num_workers=4)

# Define Patch Embedding with optional convolutional layers
class PatchEmbedding(nn.Module):
    def __init__(self, img_size=32, patch_size=4, in_channels=3, embed_dim=256):
        super(PatchEmbedding, self).__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        self.embed_dim = embed_dim
        self.conv_layers = nn.Sequential(
            nn.Conv2d(in_channels, embed_dim // 2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(embed_dim // 2),
            nn.ReLU(),
            nn.Conv2d(embed_dim // 2, embed_dim, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(embed_dim),
            nn.ReLU(),
        )
        self.proj = nn.Conv2d(embed_dim, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        x = self.conv_layers(x)
        x = self.proj(x)  # Shape: [batch_size, embed_dim, num_patches_root, num_patches_root]
        x = x.flatten(2)  # Shape: [batch_size, embed_dim, num_patches]
        x = x.transpose(1, 2)  # Shape: [batch_size, num_patches, embed_dim]
        return x

# Sequential Attention Block
class SequentialAttentionBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super(SequentialAttentionBlock, self).__init__()
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
        self.norm = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x shape: [seq_length, batch_size, embed_dim]
        seq_length = x.size(0)
        attn_mask = torch.triu(torch.ones(seq_length, seq_length), diagonal=1).bool().to(x.device)
        attn_output, _ = self.attention(x, x, x, attn_mask=attn_mask)
        x = self.norm(x + attn_output)
        return self.dropout(x)

# Cyclic Attention Block with CRF
class CyclicAttentionBlockCRF(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super(CyclicAttentionBlockCRF, self).__init__()
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
        self.norm = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)
        self.cyclic_operator = nn.Linear(embed_dim, embed_dim, bias=False)

    def forward(self, x):
        attn_output, _ = self.attention(x, x, x)
        x = self.norm(x + attn_output)
        cyclic_term = self.cyclic_alignment(attn_output)
        x = self.norm(x + cyclic_term)
        return self.dropout(x)

    def cyclic_alignment(self, attn_output):
        cyclic_term = self.cyclic_operator(attn_output)
        cyclic_term = torch.roll(cyclic_term, shifts=1, dims=0)
        return cyclic_term

# Combined Transformer Block with additional multi-headed self-attention and sequential attention
class CombinedTransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1, dropconnect_p=0.5):
        super(CombinedTransformerBlock, self).__init__()
        self.initial_attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropconnect_p)
        self.norm0 = nn.LayerNorm(embed_dim)

        self.attention1 = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropconnect_p)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.dropconnect = nn.Dropout(dropconnect_p)
        self.cyclic_attention = CyclicAttentionBlockCRF(embed_dim, num_heads, dropout)
        self.sequential_attention = SequentialAttentionBlock(embed_dim, num_heads, dropout)
        self.attention2 = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropconnect_p)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.ff = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embed_dim)
        )
        self.norm3 = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        attn_output, _ = self.initial_attention(x, x, x)
        x = self.norm0(x + attn_output)

        attn_output, _ = self.attention1(x, x, x)
        x = self.norm1(x + attn_output)
        x = self.dropconnect(x)
        x = self.cyclic_attention(x)
        x = self.sequential_attention(x)
        attn_output, _ = self.attention2(x, x, x)
        x = self.norm2(x + attn_output)
        ff_output = self.ff(x)
        x = self.norm3(x + self.dropout(ff_output))
        return x

# Decoder Block
class DecoderBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
        super(DecoderBlock, self).__init__()
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.cyclic_attention = CyclicAttentionBlockCRF(embed_dim, num_heads, dropout)
        self.ff = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embed_dim)
        )
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, encoder_output):
        attn_output, _ = self.attention(x, encoder_output, encoder_output)
        x = self.norm1(x + attn_output)
        x = self.cyclic_attention(x)
        ff_output = self.ff(x)
        x = self.norm2(x + self.dropout(ff_output))
        return x

# Custom Transformer Model with increased depth and learnable positional encodings
class CustomTransformer(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim, num_classes, num_layers=6, dropconnect_p=0.5):
        super(CustomTransformer, self).__init__()
        self.embed_dim = embed_dim
        self.num_patches = (32 // 4) ** 2  # Assuming patch_size=4
        self.patch_embedding = PatchEmbedding(embed_dim=embed_dim)
        self.positional_encoding = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))
        nn.init.trunc_normal_(self.positional_encoding, std=0.02)

        # Create multiple encoder blocks
        self.encoder_blocks = nn.ModuleList([
            CombinedTransformerBlock(embed_dim, num_heads, ff_dim, dropconnect_p=dropconnect_p)
            for _ in range(num_layers)
        ])

        # Create multiple decoder blocks
        self.decoder_blocks = nn.ModuleList([
            DecoderBlock(embed_dim, num_heads, ff_dim)
            for _ in range(num_layers)
        ])

        self.fc = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        x = self.patch_embedding(x)  # Shape: [batch_size, num_patches, embed_dim]
        x += self.positional_encoding
        x = x.transpose(0, 1)  # Shape: [num_patches, batch_size, embed_dim]

        encoder_output = x
        for encoder in self.encoder_blocks:
            encoder_output = encoder(encoder_output)

        decoder_output = encoder_output
        for decoder in self.decoder_blocks:
            decoder_output = decoder(decoder_output, encoder_output)

        decoder_output = decoder_output.mean(dim=0)  # Shape: [batch_size, embed_dim]
        logits = self.fc(decoder_output)
        return logits

# Initialize two instances of the model for 'model_epoch_125' and 'best_model'
embed_dim = 512
num_heads = 32
ff_dim = 1024
num_classes = 10
num_layers = 10  # Ensure it matches the architecture

model_epoch_125 = CustomTransformer(embed_dim, num_heads, ff_dim, num_classes, num_layers=num_layers).to(device)
model_best = CustomTransformer(embed_dim, num_heads, ff_dim, num_classes, num_layers=num_layers).to(device)

# Download the models from Hugging Face Hub
from huggingface_hub import hf_hub_download

# Paths where the models will be saved
model_epoch_125_path = hf_hub_download(repo_id=f"{username}/{model_name_epoch_125}", filename="model_epoch_125.pth")
model_best_path = hf_hub_download(repo_id=f"{username}/{model_name_best}", filename="model_best.pth")

# Load the saved models from Hugging Face Hub
model_epoch_125.load_state_dict(torch.load(model_epoch_125_path, map_location=device))
model_best.load_state_dict(torch.load(model_best_path, map_location=device))

# Set both models to evaluation mode
model_epoch_125.eval()
model_best.eval()

# Prepare the feature and label arrays
test_preds_epoch_125 = []
test_preds_best = []
test_labels = []

with torch.no_grad():
    for images_test, labels_test in test_loader:
        images_test = images_test.to(device)

        # Get predictions from model_epoch_125
        logits_epoch_125 = model_epoch_125(images_test)
        probs_epoch_125 = F.softmax(logits_epoch_125, dim=1).cpu().numpy()  # Convert to probabilities

        # Get predictions from model_best
        logits_best = model_best(images_test)
        probs_best = F.softmax(logits_best, dim=1).cpu().numpy()  # Convert to probabilities

        # Store predictions and labels
        test_preds_epoch_125.extend(probs_epoch_125)
        test_preds_best.extend(probs_best)
        test_labels.extend(labels_test.numpy())

# Convert predictions to NumPy arrays
test_preds_epoch_125 = np.array(test_preds_epoch_125)
test_preds_best = np.array(test_preds_best)
test_labels = np.array(test_labels)

# Stack the predictions from both models to create meta-features
meta_features = np.hstack((test_preds_epoch_125, test_preds_best))  # Shape: (num_samples, 20)

# Split the data for training and validation of the XGBoost meta-learner
X_train, X_val, y_train, y_val = train_test_split(meta_features, test_labels, test_size=0.2, random_state=42)

# Train an XGBoost classifier as a meta-learner
xgb_model = xgb.XGBClassifier(
    objective='multi:softmax',
    num_class=10,
    eval_metric='mlogloss',
    use_label_encoder=False
)

xgb_model.fit(X_train, y_train)

# Validate the XGBoost model
y_pred_val = xgb_model.predict(X_val)
val_accuracy = accuracy_score(y_val, y_pred_val)
print(f'Validation Accuracy of Meta-learner: {val_accuracy * 100:.2f}%')

# Test the XGBoost model on the entire test set
y_pred_test = xgb_model.predict(meta_features)
test_accuracy = accuracy_score(test_labels, y_pred_test)
print(f'Test Accuracy of Meta-learner: {test_accuracy * 100:.2f}%')

# Plot the confusion matrix for the test set predictions
cm = confusion_matrix(test_labels, y_pred_test)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=cifar10_test.classes)
disp.plot(cmap=plt.cm.Blues)

# Rotate the x-axis labels to prevent overlapping
plt.xticks(rotation=45, ha='right')
plt.title('Confusion Matrix for XGBoost Meta-learner on CIFAR-10 Test Set')
plt.savefig(os.path.join(save_dir, 'xgboost_meta_confusion_matrix.png'))
plt.show()

# Save the XGBoost model
xgb_model.save_model(os.path.join(save_dir, 'xgboost_meta_learner.json'))
print('Meta-learner model saved.')