Upload test1/Algo_CIFAR_10_MobileViTv3.py with huggingface_hub

test1/Algo_CIFAR_10_MobileViTv3.py

@@ -0,0 +1,447 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as transforms
+from codecarbon import EmissionsTracker
+from carbontracker.tracker import CarbonTracker
+from fvcore.nn import FlopCountAnalysis
+from sklearn.metrics import precision_recall_fscore_support, accuracy_score
+from einops import rearrange
+from tqdm import tqdm
+import pandas as pd
+import numpy as np
+import pickle
+import os
+import time
+import logging
+import warnings
+import gc
+
+# ==========================================
+# 1. MOBILEVITV3 ARCHITECTURE DEFINITION
+# ==========================================
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.SiLU()
+    )
+
+def conv_nxn_bn(inp, oup, kernel_size=3, stride=1):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, kernel_size, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.SiLU()
+    )
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+# --- SPEED FIX 1: Hardware-Fused Attention Kernel ---
+class Attention(nn.Module):
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.dropout_rate = dropout 
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h=self.heads), qkv)
+
+        # PyTorch native SDPA triggers FlashAttention on supported hardware
+        out = F.scaled_dot_product_attention(
+            q, k, v, 
+            dropout_p=self.dropout_rate if self.training else 0.0
+        )
+
+        out = rearrange(out, 'b p h n d -> b p n (h d)')
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
+                PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
+            ]))
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return x
+
+class MV2Block(nn.Module):
+    def __init__(self, inp, oup, stride=1, expansion=4):
+        super().__init__()
+        self.stride = stride
+        hidden_dim = int(inp * expansion)
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        if expansion == 1:
+            self.conv = nn.Sequential(
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.SiLU(),
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.SiLU(),
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.SiLU(),
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+class MobileViTBlock(nn.Module):
+    def __init__(self, dim, depth, channel, kernel_size, patch_size, mlp_dim, dropout=0.):
+        super().__init__()
+        self.ph, self.pw = patch_size
+
+        self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
+        self.conv2 = conv_1x1_bn(channel, dim)
+
+        self.transformer = Transformer(dim, depth, 1, 32, mlp_dim, dropout)
+
+        self.conv3 = conv_1x1_bn(dim, channel)
+        self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)
+    
+    def forward(self, x):
+        y = x.clone()
+
+        x = self.conv1(x)
+        x = self.conv2(x)
+        
+        _, _, h, w = x.shape
+        pad_h = (self.ph - h % self.ph) % self.ph
+        pad_w = (self.pw - w % self.pw) % self.pw
+        
+        if pad_h > 0 or pad_w > 0:
+            x = nn.functional.pad(x, (0, pad_w, 0, pad_h))
+        
+        _, _, h_pad, w_pad = x.shape
+        x = rearrange(x, 'b d (h ph) (w pw) -> b (ph pw) (h w) d', ph=self.ph, pw=self.pw)
+        x = self.transformer(x)
+        x = rearrange(x, 'b (ph pw) (h w) d -> b d (h ph) (w pw)', h=h_pad//self.ph, w=w_pad//self.pw, ph=self.ph, pw=self.pw)
+
+        if pad_h > 0 or pad_w > 0:
+            x = x[:, :, :h, :w]
+
+        x = self.conv3(x)
+        x = torch.cat((x, y), 1)
+        x = self.conv4(x)
+        return x
+
+class MobileViTv3_Small(nn.Module):
+    def __init__(self, image_size=(224, 224), num_classes=10):
+        super().__init__()
+        ih, iw = image_size
+        ph, pw = 2, 2
+        
+        dims = [144, 192, 240]
+        channels = [16, 32, 64, 64, 96, 96, 128, 128, 160, 160, 640]
+        
+        self.conv1 = conv_nxn_bn(3, channels[0], stride=2)
+
+        self.mv2 = nn.ModuleList([])
+        self.mv2.append(MV2Block(channels[0], channels[1], 1, 4))
+        self.mv2.append(MV2Block(channels[1], channels[2], 2, 4))
+        self.mv2.append(MV2Block(channels[2], channels[3], 1, 4))
+        self.mv2.append(MV2Block(channels[3], channels[4], 2, 4))
+        
+        self.mvit = nn.ModuleList([])
+        self.mvit.append(MobileViTBlock(dims[0], 2, channels[5], 3, (ph, pw), int(dims[0]*2)))
+        
+        self.mv2_2 = nn.ModuleList([])
+        self.mv2_2.append(MV2Block(channels[5], channels[6], 2, 4))
+        
+        self.mvit_2 = nn.ModuleList([])
+        self.mvit_2.append(MobileViTBlock(dims[1], 4, channels[7], 3, (ph, pw), int(dims[1]*2)))
+        
+        self.mv2_3 = nn.ModuleList([])
+        self.mv2_3.append(MV2Block(channels[7], channels[8], 2, 4))
+        
+        self.mvit_3 = nn.ModuleList([])
+        self.mvit_3.append(MobileViTBlock(dims[2], 3, channels[9], 3, (ph, pw), int(dims[2]*2)))
+        
+        self.conv2 = conv_1x1_bn(channels[9], channels[10])
+        self.pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(channels[10], num_classes)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        for conv in self.mv2: x = conv(x)
+        for m in self.mvit: x = m(x)
+        for conv in self.mv2_2: x = conv(x)
+        for m in self.mvit_2: x = m(x)
+        for conv in self.mv2_3: x = conv(x)
+        for m in self.mvit_3: x = m(x)
+        x = self.conv2(x)
+        x = self.pool(x).view(-1, x.shape[1])
+        return self.fc(x)
+
+# ==========================================
+# 2. EDEN EXPERIMENT & PROFILING
+# ==========================================
+
+warnings.filterwarnings("ignore", category=UserWarning) 
+logging.getLogger("codecarbon").setLevel(logging.ERROR) 
+
+# --- Configurations ---
+DATA_DIR = r"C:\Users\shanm\Dataset Download\cifar-10-batches-py"
+LOG_FILE = "eden_optimized_cifar10_custom_mobilevitv3.csv"
+MODEL_SAVE_PATH = "eden_optimized_custom_mobilevitv3_cifar10.pth"
+
+BATCH_SIZE = 32
+ACCUMULATION_STEPS = 4
+LEARNING_RATE = 1e-3
+NUM_EPOCHS = 50
+L1_LAMBDA = 1e-5        
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# --- SPEED FIX 2: Barebones CPU DataLoader ---
+class CIFAR10Binary(Dataset):
+    def __init__(self, root, train=True):
+        self.data = []
+        self.labels = []
+        
+        if train:
+            for i in range(1, 6):
+                file_path = os.path.join(root, f'data_batch_{i}')
+                with open(file_path, 'rb') as f:
+                    entry = pickle.load(f, encoding='latin1')
+                    self.data.append(entry['data'])
+                    self.labels.extend(entry['labels'])
+            self.data = np.vstack(self.data).reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
+        else:
+            file_path = os.path.join(root, 'test_batch')
+            with open(file_path, 'rb') as f:
+                entry = pickle.load(f, encoding='latin1')
+                self.data = entry['data'].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
+                self.labels = entry['labels']
+
+    def __len__(self): return len(self.data)
+    
+    def __getitem__(self, idx):
+        img, target = self.data[idx], self.labels[idx]
+        # Instantly convert raw numpy to tensor (no PIL conversion)
+        img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
+        return img, target
+
+def run_experiment():
+    # --- SPEED FIX 3: Enable CuDNN Benchmark ---
+    torch.backends.cudnn.benchmark = True
+    torch.cuda.empty_cache()
+    gc.collect()
+    
+    print("\n[EDEN Setup] Loading Custom MobileViTv3-Small...")
+    model = MobileViTv3_Small(image_size=(224, 224), num_classes=10)
+    model = model.to(DEVICE)
+    
+    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+    dummy_input = torch.randn(1, 3, 224, 224).to(DEVICE)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        total_flops = FlopCountAnalysis(model, dummy_input).total()
+    total_params = sum(p.numel() for p in model.parameters())
+    
+    # Load dataset with zero CPU preprocessing overhead
+    train_set = CIFAR10Binary(root=DATA_DIR, train=True)
+    loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)
+
+    criterion = nn.CrossEntropyLoss()
+    scaler = torch.cuda.amp.GradScaler() 
+    
+    cc_tracker = EmissionsTracker(measure_power_secs=1, save_to_file=False)
+    ct_tracker = CarbonTracker(epochs=NUM_EPOCHS, monitor_epochs=NUM_EPOCHS, update_interval=1)
+    
+    cc_tracker.start()
+    all_logs = []
+    total_iterations_counter = 0
+    session_start_time = time.time()
+
+    prev_cum_gpu_j, prev_cum_cpu_j, prev_cum_ram_j = 0.0, 0.0, 0.0
+
+    # Define the normalizer once, push it directly to the GPU
+    normalizer = transforms.Normalize(
+        (0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)
+    ).to(DEVICE)
+
+    print(f"\nEDEN PROFILING STARTED | DEVICE: {torch.cuda.get_device_name(0)}")
+    print(f"Dataset: CIFAR-10 (RAM Cached, GPU Resized) | Architecture: Custom MobileViTv3-Small")
+    print(f"Params: {total_params:,} | FLOPs: {total_flops:.2e}\n")
+
+    for epoch in range(NUM_EPOCHS):
+        ct_tracker.epoch_start()
+        torch.cuda.reset_peak_memory_stats() 
+        epoch_start_time = time.time()
+        model.train()
+        
+        running_loss = 0.0
+        all_preds, all_labels = [], []
+        epoch_grad_norms = []
+        
+        optimizer.zero_grad()
+        pbar = tqdm(loader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS}", unit="batch", leave=False)
+        
+        for i, (images, labels) in enumerate(pbar):
+            # --- SPEED FIX 4: GPU-Accelerated Resizing and Normalization ---
+            images, labels = images.to(DEVICE), labels.to(DEVICE)
+            
+            # Resize the entire batch simultaneously on the CUDA cores
+            images = F.interpolate(
+                images, size=(224, 224), mode='bilinear', align_corners=False
+            )
+            images = normalizer(images)
+            # ---------------------------------------------------------------
+
+            with torch.cuda.amp.autocast():
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                
+                # Active Sparse Training (L1 Penalty)
+                l1_penalty = sum(p.abs().sum() for p in model.parameters())
+                total_loss = loss + (L1_LAMBDA * l1_penalty)
+                scaled_loss = total_loss / ACCUMULATION_STEPS
+
+            scaler.scale(scaled_loss).backward()
+            
+            grad_norm = 0.0
+            for p in model.parameters():
+                if p.grad is not None:
+                    grad_norm += p.grad.data.norm(2).item() ** 2
+            epoch_grad_norms.append(grad_norm ** 0.5)
+
+            if (i + 1) % ACCUMULATION_STEPS == 0:
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
+
+            running_loss += loss.item() * ACCUMULATION_STEPS
+            
+            _, preds = torch.max(outputs, 1)
+            all_preds.extend(preds.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+            total_iterations_counter += 1
+            
+            pbar.set_postfix(loss=f"{(loss.item()*ACCUMULATION_STEPS):.4f}")
+
+        # --- A. Evaluation ---
+        ct_tracker.epoch_end()
+        epoch_end_time = time.time()
+        epoch_duration = epoch_end_time - epoch_start_time
+        avg_it_per_sec = len(loader) / epoch_duration
+        
+        acc = accuracy_score(all_labels, all_preds)
+        p, r, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='macro', zero_division=0)
+        
+        model.eval()
+        with torch.no_grad():
+            sample_img = torch.randn(1, 3, 224, 224).to(DEVICE)
+            _ = model(sample_img) 
+            torch.cuda.synchronize()
+            
+            starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+            starter.record()
+            _ = model(sample_img) 
+            ender.record()
+            torch.cuda.synchronize()
+            lat_ms = starter.elapsed_time(ender)
+
+        # --- B. Energy & Power Calculations ---
+        emissions_data = cc_tracker._prepare_emissions_data()
+        
+        cum_gpu_j = emissions_data.gpu_energy * 3.6e6
+        cum_cpu_j = emissions_data.cpu_energy * 3.6e6
+        cum_ram_j = emissions_data.ram_energy * 3.6e6
+        cum_total_j = cum_gpu_j + cum_cpu_j + cum_ram_j
+        
+        epoch_gpu_j = cum_gpu_j - prev_cum_gpu_j
+        epoch_cpu_j = cum_cpu_j - prev_cum_cpu_j
+        epoch_ram_j = cum_ram_j - prev_cum_ram_j
+        epoch_total_j = epoch_gpu_j + epoch_cpu_j + epoch_ram_j
+        
+        prev_cum_gpu_j, prev_cum_cpu_j, prev_cum_ram_j = cum_gpu_j, cum_cpu_j, cum_ram_j
+
+        avg_gpu_w = epoch_gpu_j / epoch_duration if epoch_duration > 0 else 0
+        avg_cpu_w = epoch_cpu_j / epoch_duration if epoch_duration > 0 else 0
+        avg_ram_w = epoch_ram_j / epoch_duration if epoch_duration > 0 else 0
+
+        vram_peak = torch.cuda.max_memory_allocated(DEVICE) / (1024**3)
+        
+        # --- C. Terminal Update ---
+        print(f"Epoch {epoch+1} Summary:")
+        print(f" > Acc: {acc:.4f} | F1: {f1:.4f} | Loss: {running_loss/len(loader):.4f}")
+        print(f" > Epoch Energy: {epoch_total_j:.1f}J")
+        print(f" > Avg Power: GPU {avg_gpu_w:.1f}W | VRAM: {vram_peak:.2f}GB | Latency: {lat_ms:.2f}ms")
+        print("-" * 65)
+
+        # --- D. Unified Verified CSV Logging ---
+        log_entry = {
+            "epoch": epoch + 1,
+            "loss": running_loss / len(loader),
+            "accuracy": acc, "f1_score": f1, "precision": p, "recall": r,
+            "epoch_energy_gpu_j": epoch_gpu_j, "epoch_energy_cpu_j": epoch_cpu_j,
+            "epoch_energy_ram_j": epoch_ram_j, "epoch_total_energy_j": epoch_total_j,
+            "cumulative_total_energy_j": cum_total_j, "carbon_emissions_kg": emissions_data.emissions,
+            "avg_power_gpu_w": avg_gpu_w, "avg_power_cpu_w": avg_cpu_w, "avg_power_ram_w": avg_ram_w,
+            "vram_peak_gb": vram_peak, "latency_ms": lat_ms, "avg_grad_norm": np.mean(epoch_grad_norms),
+            "it_per_sec": avg_it_per_sec, "total_iterations": total_iterations_counter,
+            "epoch_duration_sec": epoch_duration, "cumulative_time_sec": time.time() - session_start_time
+        }
+        all_logs.append(log_entry)
+        pd.DataFrame(all_logs).to_csv(LOG_FILE, index=False)
+
+    cc_tracker.stop()
+    
+    # --- E. Save Optimized Model ---
+    torch.save(model.state_dict(), MODEL_SAVE_PATH)
+    print(f"\n[FINISH] Verified Optimization Complete.")
+
+if __name__ == "__main__":
+    run_experiment()