Shanmuk4622 commited on
Commit
c09c98b
·
verified ·
1 Parent(s): 2305b52

Upload test1/Algo_ImageNet_mobilevit3.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. test1/Algo_ImageNet_mobilevit3.py +416 -0
test1/Algo_ImageNet_mobilevit3.py ADDED
@@ -0,0 +1,416 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.optim as optim
4
+ import torch.nn.functional as F
5
+ from torch.utils.data import DataLoader
6
+ from torchvision.datasets import ImageFolder
7
+ import torchvision.transforms as transforms
8
+ from codecarbon import EmissionsTracker
9
+ from carbontracker.tracker import CarbonTracker
10
+ from fvcore.nn import FlopCountAnalysis
11
+ from sklearn.metrics import precision_recall_fscore_support, accuracy_score
12
+ from einops import rearrange
13
+ from tqdm import tqdm
14
+ import pandas as pd
15
+ import numpy as np
16
+ import os
17
+ import time
18
+ import logging
19
+ import warnings
20
+ import gc
21
+
22
+ # --- Environment & Logging Optimization ---
23
+ warnings.filterwarnings("ignore", category=UserWarning)
24
+ # Hard-mute CodeCarbon terminal spam
25
+ logging.getLogger("codecarbon").setLevel(logging.CRITICAL)
26
+ logging.getLogger("codecarbon").disabled = True
27
+
28
+ # --- Configurations ---
29
+ DATA_DIR = r"C:\Users\shanm\Dataset Download\custom image net"
30
+ LOG_FILE = "eden_optimized_custom_imagenet_mobilevitv3.csv"
31
+ MODEL_SAVE_PATH = "eden_optimized_custom_mobilevitv3_imagenet.pth"
32
+
33
+ BATCH_SIZE = 32
34
+ ACCUMULATION_STEPS = 4
35
+ LEARNING_RATE = 1e-3
36
+ NUM_EPOCHS = 30
37
+ L1_LAMBDA = 1e-5
38
+ NUM_CLASSES = 300 # Matched to your 300 custom folders
39
+
40
+ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
41
+
42
+ # ==========================================
43
+ # 1. MOBILEVITV3 ARCHITECTURE DEFINITION
44
+ # ==========================================
45
+
46
+ def conv_1x1_bn(inp, oup):
47
+ return nn.Sequential(
48
+ nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
49
+ nn.BatchNorm2d(oup),
50
+ nn.SiLU()
51
+ )
52
+
53
+ def conv_nxn_bn(inp, oup, kernel_size=3, stride=1):
54
+ return nn.Sequential(
55
+ nn.Conv2d(inp, oup, kernel_size, stride, 1, bias=False),
56
+ nn.BatchNorm2d(oup),
57
+ nn.SiLU()
58
+ )
59
+
60
+ class PreNorm(nn.Module):
61
+ def __init__(self, dim, fn):
62
+ super().__init__()
63
+ self.norm = nn.LayerNorm(dim)
64
+ self.fn = fn
65
+ def forward(self, x, **kwargs):
66
+ return self.fn(self.norm(x), **kwargs)
67
+
68
+ class FeedForward(nn.Module):
69
+ def __init__(self, dim, hidden_dim, dropout=0.):
70
+ super().__init__()
71
+ self.net = nn.Sequential(
72
+ nn.Linear(dim, hidden_dim),
73
+ nn.SiLU(),
74
+ nn.Dropout(dropout),
75
+ nn.Linear(hidden_dim, dim),
76
+ nn.Dropout(dropout)
77
+ )
78
+ def forward(self, x):
79
+ return self.net(x)
80
+
81
+ # Hardware-Fused Attention Kernel for Maximum Speed
82
+ class Attention(nn.Module):
83
+ def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
84
+ super().__init__()
85
+ inner_dim = dim_head * heads
86
+ project_out = not (heads == 1 and dim_head == dim)
87
+
88
+ self.heads = heads
89
+ self.dropout_rate = dropout
90
+
91
+ self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
92
+
93
+ self.to_out = nn.Sequential(
94
+ nn.Linear(inner_dim, dim),
95
+ nn.Dropout(dropout)
96
+ ) if project_out else nn.Identity()
97
+
98
+ def forward(self, x):
99
+ qkv = self.to_qkv(x).chunk(3, dim=-1)
100
+ q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h=self.heads), qkv)
101
+
102
+ # PyTorch native SDPA triggers FlashAttention
103
+ out = F.scaled_dot_product_attention(
104
+ q, k, v,
105
+ dropout_p=self.dropout_rate if self.training else 0.0
106
+ )
107
+
108
+ out = rearrange(out, 'b p h n d -> b p n (h d)')
109
+ return self.to_out(out)
110
+
111
+ class Transformer(nn.Module):
112
+ def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
113
+ super().__init__()
114
+ self.layers = nn.ModuleList([])
115
+ for _ in range(depth):
116
+ self.layers.append(nn.ModuleList([
117
+ PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
118
+ PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
119
+ ]))
120
+ def forward(self, x):
121
+ for attn, ff in self.layers:
122
+ x = attn(x) + x
123
+ x = ff(x) + x
124
+ return x
125
+
126
+ class MV2Block(nn.Module):
127
+ def __init__(self, inp, oup, stride=1, expansion=4):
128
+ super().__init__()
129
+ self.stride = stride
130
+ hidden_dim = int(inp * expansion)
131
+ self.use_res_connect = self.stride == 1 and inp == oup
132
+
133
+ if expansion == 1:
134
+ self.conv = nn.Sequential(
135
+ nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
136
+ nn.BatchNorm2d(hidden_dim),
137
+ nn.SiLU(),
138
+ nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
139
+ nn.BatchNorm2d(oup),
140
+ )
141
+ else:
142
+ self.conv = nn.Sequential(
143
+ nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
144
+ nn.BatchNorm2d(hidden_dim),
145
+ nn.SiLU(),
146
+ nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
147
+ nn.BatchNorm2d(hidden_dim),
148
+ nn.SiLU(),
149
+ nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
150
+ nn.BatchNorm2d(oup),
151
+ )
152
+
153
+ def forward(self, x):
154
+ if self.use_res_connect:
155
+ return x + self.conv(x)
156
+ else:
157
+ return self.conv(x)
158
+
159
+ class MobileViTBlock(nn.Module):
160
+ def __init__(self, dim, depth, channel, kernel_size, patch_size, mlp_dim, dropout=0.):
161
+ super().__init__()
162
+ self.ph, self.pw = patch_size
163
+
164
+ self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
165
+ self.conv2 = conv_1x1_bn(channel, dim)
166
+
167
+ self.transformer = Transformer(dim, depth, 1, 32, mlp_dim, dropout)
168
+
169
+ self.conv3 = conv_1x1_bn(dim, channel)
170
+ self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)
171
+
172
+ def forward(self, x):
173
+ y = x.clone()
174
+
175
+ x = self.conv1(x)
176
+ x = self.conv2(x)
177
+
178
+ _, _, h, w = x.shape
179
+ pad_h = (self.ph - h % self.ph) % self.ph
180
+ pad_w = (self.pw - w % self.pw) % self.pw
181
+
182
+ if pad_h > 0 or pad_w > 0:
183
+ x = nn.functional.pad(x, (0, pad_w, 0, pad_h))
184
+
185
+ _, _, h_pad, w_pad = x.shape
186
+ x = rearrange(x, 'b d (h ph) (w pw) -> b (ph pw) (h w) d', ph=self.ph, pw=self.pw)
187
+ x = self.transformer(x)
188
+ 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)
189
+
190
+ if pad_h > 0 or pad_w > 0:
191
+ x = x[:, :, :h, :w]
192
+
193
+ x = self.conv3(x)
194
+ x = torch.cat((x, y), 1)
195
+ x = self.conv4(x)
196
+ return x
197
+
198
+ class MobileViTv3_Small(nn.Module):
199
+ def __init__(self, image_size=(224, 224), num_classes=300): # Updated for 300 Classes
200
+ super().__init__()
201
+ ih, iw = image_size
202
+ ph, pw = 2, 2
203
+
204
+ dims = [144, 192, 240]
205
+ channels = [16, 32, 64, 64, 96, 96, 128, 128, 160, 160, 640]
206
+
207
+ self.conv1 = conv_nxn_bn(3, channels[0], stride=2)
208
+
209
+ self.mv2 = nn.ModuleList([])
210
+ self.mv2.append(MV2Block(channels[0], channels[1], 1, 4))
211
+ self.mv2.append(MV2Block(channels[1], channels[2], 2, 4))
212
+ self.mv2.append(MV2Block(channels[2], channels[3], 1, 4))
213
+ self.mv2.append(MV2Block(channels[3], channels[4], 2, 4))
214
+
215
+ self.mvit = nn.ModuleList([])
216
+ self.mvit.append(MobileViTBlock(dims[0], 2, channels[5], 3, (ph, pw), int(dims[0]*2)))
217
+
218
+ self.mv2_2 = nn.ModuleList([])
219
+ self.mv2_2.append(MV2Block(channels[5], channels[6], 2, 4))
220
+
221
+ self.mvit_2 = nn.ModuleList([])
222
+ self.mvit_2.append(MobileViTBlock(dims[1], 4, channels[7], 3, (ph, pw), int(dims[1]*2)))
223
+
224
+ self.mv2_3 = nn.ModuleList([])
225
+ self.mv2_3.append(MV2Block(channels[7], channels[8], 2, 4))
226
+
227
+ self.mvit_3 = nn.ModuleList([])
228
+ self.mvit_3.append(MobileViTBlock(dims[2], 3, channels[9], 3, (ph, pw), int(dims[2]*2)))
229
+
230
+ self.conv2 = conv_1x1_bn(channels[9], channels[10])
231
+ self.pool = nn.AdaptiveAvgPool2d((1, 1))
232
+ self.fc = nn.Linear(channels[10], num_classes)
233
+
234
+ def forward(self, x):
235
+ x = self.conv1(x)
236
+ for conv in self.mv2: x = conv(x)
237
+ for m in self.mvit: x = m(x)
238
+ for conv in self.mv2_2: x = conv(x)
239
+ for m in self.mvit_2: x = m(x)
240
+ for conv in self.mv2_3: x = conv(x)
241
+ for m in self.mvit_3: x = m(x)
242
+ x = self.conv2(x)
243
+ x = self.pool(x).view(-1, x.shape[1])
244
+ return self.fc(x)
245
+
246
+ # ==========================================
247
+ # 2. EDEN EXPERIMENT & PROFILING
248
+ # ==========================================
249
+
250
+ def run_experiment():
251
+ torch.backends.cudnn.benchmark = True
252
+ torch.cuda.empty_cache()
253
+ gc.collect()
254
+
255
+ # Initialize custom architecture (Training from scratch)
256
+ model = MobileViTv3_Small(image_size=(224, 224), num_classes=NUM_CLASSES)
257
+ model = model.to(DEVICE)
258
+
259
+ optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
260
+
261
+ dummy_input = torch.randn(1, 3, 224, 224).to(DEVICE)
262
+ with warnings.catch_warnings():
263
+ warnings.simplefilter("ignore")
264
+ total_flops = FlopCountAnalysis(model, dummy_input).total()
265
+ total_params = sum(p.numel() for p in model.parameters())
266
+
267
+ # --- Data Loading: Minimal CPU Processing ---
268
+ cpu_transform = transforms.Compose([
269
+ transforms.Resize((224, 224)),
270
+ transforms.ToTensor()
271
+ ])
272
+
273
+ # Directly loads from the 300 custom class folders
274
+ train_set = ImageFolder(root=DATA_DIR, transform=cpu_transform)
275
+ loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)
276
+
277
+ criterion = nn.CrossEntropyLoss()
278
+ scaler = torch.cuda.amp.GradScaler()
279
+
280
+ # --- 3. Profiling Initialization (SILENCED) ---
281
+ cc_tracker = EmissionsTracker(measure_power_secs=1, save_to_file=False, log_level="critical")
282
+ ct_tracker = CarbonTracker(epochs=NUM_EPOCHS, monitor_epochs=NUM_EPOCHS, update_interval=1)
283
+
284
+ cc_tracker.start()
285
+ all_logs = []
286
+ total_iterations_counter = 0
287
+ session_start_time = time.time()
288
+
289
+ prev_cum_gpu_j, prev_cum_cpu_j, prev_cum_ram_j = 0.0, 0.0, 0.0
290
+
291
+ # ImageNet Standard Normalizer pushed directly to the GPU
292
+ gpu_normalizer = transforms.Normalize(
293
+ mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
294
+ ).to(DEVICE)
295
+
296
+ print(f"\n[EDEN PROFILING STARTED] | Model: Custom MobileViTv3-Small | Classes: {NUM_CLASSES}")
297
+ print(f"Dataset: Custom ImageNet ({len(train_set)} images) | Saving quietly to CSV...\n")
298
+
299
+ for epoch in range(NUM_EPOCHS):
300
+ ct_tracker.epoch_start()
301
+ torch.cuda.reset_peak_memory_stats()
302
+ epoch_start_time = time.time()
303
+ model.train()
304
+
305
+ running_loss = 0.0
306
+ all_preds, all_labels = [], []
307
+ epoch_grad_norms = []
308
+
309
+ optimizer.zero_grad()
310
+ pbar = tqdm(loader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS}", unit="batch", leave=False)
311
+
312
+ for i, (images, labels) in enumerate(pbar):
313
+ # GPU-Accelerated Normalization
314
+ images, labels = images.to(DEVICE), labels.to(DEVICE)
315
+ images = gpu_normalizer(images)
316
+
317
+ with torch.cuda.amp.autocast():
318
+ outputs = model(images)
319
+ loss = criterion(outputs, labels)
320
+
321
+ # Active Sparse Training (L1 Penalty) applied natively
322
+ l1_penalty = sum(p.abs().sum() for p in model.parameters())
323
+ total_loss = loss + (L1_LAMBDA * l1_penalty)
324
+ scaled_loss = total_loss / ACCUMULATION_STEPS
325
+
326
+ scaler.scale(scaled_loss).backward()
327
+
328
+ grad_norm = 0.0
329
+ for p in model.parameters():
330
+ if p.grad is not None:
331
+ grad_norm += p.grad.data.norm(2).item() ** 2
332
+ epoch_grad_norms.append(grad_norm ** 0.5)
333
+
334
+ if (i + 1) % ACCUMULATION_STEPS == 0:
335
+ scaler.step(optimizer)
336
+ scaler.update()
337
+ optimizer.zero_grad()
338
+
339
+ running_loss += loss.item() * ACCUMULATION_STEPS
340
+
341
+ _, preds = torch.max(outputs, 1)
342
+ all_preds.extend(preds.cpu().numpy())
343
+ all_labels.extend(labels.cpu().numpy())
344
+ total_iterations_counter += 1
345
+
346
+ pbar.set_postfix(loss=f"{(loss.item()*ACCUMULATION_STEPS):.4f}")
347
+
348
+ # --- A. Evaluation ---
349
+ ct_tracker.epoch_end()
350
+ epoch_end_time = time.time()
351
+ epoch_duration = epoch_end_time - epoch_start_time
352
+ avg_it_per_sec = len(loader) / epoch_duration
353
+
354
+ acc = accuracy_score(all_labels, all_preds)
355
+ p, r, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='macro', zero_division=0)
356
+
357
+ model.eval()
358
+ with torch.no_grad():
359
+ sample_img = torch.randn(1, 3, 224, 224).to(DEVICE)
360
+ _ = model(sample_img)
361
+ torch.cuda.synchronize()
362
+
363
+ starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
364
+ starter.record()
365
+ _ = model(sample_img)
366
+ ender.record()
367
+ torch.cuda.synchronize()
368
+ lat_ms = starter.elapsed_time(ender)
369
+
370
+ # --- B. Energy & Power Calculations ---
371
+ emissions_data = cc_tracker._prepare_emissions_data()
372
+
373
+ cum_gpu_j = emissions_data.gpu_energy * 3.6e6
374
+ cum_cpu_j = emissions_data.cpu_energy * 3.6e6
375
+ cum_ram_j = emissions_data.ram_energy * 3.6e6
376
+ cum_total_j = cum_gpu_j + cum_cpu_j + cum_ram_j
377
+
378
+ epoch_gpu_j = cum_gpu_j - prev_cum_gpu_j
379
+ epoch_cpu_j = cum_cpu_j - prev_cum_cpu_j
380
+ epoch_ram_j = cum_ram_j - prev_cum_ram_j
381
+ epoch_total_j = epoch_gpu_j + epoch_cpu_j + epoch_ram_j
382
+
383
+ prev_cum_gpu_j, prev_cum_cpu_j, prev_cum_ram_j = cum_gpu_j, cum_cpu_j, cum_ram_j
384
+
385
+ avg_gpu_w = epoch_gpu_j / epoch_duration if epoch_duration > 0 else 0
386
+ avg_cpu_w = epoch_cpu_j / epoch_duration if epoch_duration > 0 else 0
387
+ avg_ram_w = epoch_ram_j / epoch_duration if epoch_duration > 0 else 0
388
+
389
+ vram_peak = torch.cuda.max_memory_allocated(DEVICE) / (1024**3)
390
+
391
+ # --- C. Minimal Terminal Update ---
392
+ print(f"Epoch {epoch+1}/{NUM_EPOCHS} | Acc: {acc:.4f} | Loss: {running_loss/len(loader):.4f} | Energy: {epoch_total_j:.1f}J | Latency: {lat_ms:.2f}ms")
393
+
394
+ # --- D. Unified Verified CSV Logging ---
395
+ log_entry = {
396
+ "epoch": epoch + 1,
397
+ "loss": running_loss / len(loader),
398
+ "accuracy": acc, "f1_score": f1, "precision": p, "recall": r,
399
+ "epoch_energy_gpu_j": epoch_gpu_j, "epoch_energy_cpu_j": epoch_cpu_j,
400
+ "epoch_energy_ram_j": epoch_ram_j, "epoch_total_energy_j": epoch_total_j,
401
+ "cumulative_total_energy_j": cum_total_j, "carbon_emissions_kg": emissions_data.emissions,
402
+ "avg_power_gpu_w": avg_gpu_w, "avg_power_cpu_w": avg_cpu_w, "avg_power_ram_w": avg_ram_w,
403
+ "vram_peak_gb": vram_peak, "latency_ms": lat_ms, "avg_grad_norm": np.mean(epoch_grad_norms),
404
+ "it_per_sec": avg_it_per_sec, "total_iterations": total_iterations_counter,
405
+ "epoch_duration_sec": epoch_duration, "cumulative_time_sec": time.time() - session_start_time
406
+ }
407
+ all_logs.append(log_entry)
408
+ pd.DataFrame(all_logs).to_csv(LOG_FILE, index=False)
409
+
410
+ cc_tracker.stop()
411
+
412
+ torch.save(model.state_dict(), MODEL_SAVE_PATH)
413
+ print(f"\n[FINISH] Verified Optimization Complete. Model and CSV Saved.")
414
+
415
+ if __name__ == "__main__":
416
+ run_experiment()