import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from bandwidth_linear import LinearWithBandwidth class BandwidthMLP(nn.Module): """ 3x Linear(3072, 100) stacked and reshaped to (N, 3, 10, 10), then a single Conv2D(kernel=3) processes the combined output. """ def __init__(self, img_size=32, num_classes=10): super().__init__() # Input: 32x32x3 = 3072 flattened self.input_dim = img_size * img_size * 3 self.num_branches = 3 self.features_per_branch = 100 self.grid_h = self.grid_w = 10 # 10*10 = 100 # Three parallel linear layers with bandwidth self.fc1 = LinearWithBandwidth(self.input_dim, self.features_per_branch, self.grid_h) self.fc2 = LinearWithBandwidth(self.input_dim, self.features_per_branch, self.grid_h) self.fc3 = LinearWithBandwidth(self.input_dim, self.features_per_branch, self.grid_h) # Single Conv2D after stacking the 3 reshaped outputs self.conv = nn.Conv2d( in_channels=self.num_branches, out_channels=3, kernel_size=3, padding=1 ) self.bn = nn.BatchNorm2d(3) # Classifier self.pool = nn.AdaptiveAvgPool2d(1) self.classifier = nn.Sequential( #nn.Dropout(0.3), nn.Linear(300, num_classes) ) self.relu = nn.ReLU() def forward(self, x, return_intermediates=False): batch_size = x.shape[0] x_flat = x.view(batch_size, -1) # (N, 3072) # Three parallel linear+bandwidth branches out1, bw1 = self.fc1(x_flat) out2, bw2 = self.fc2(x_flat) out3, bw3 = self.fc3(x_flat) # Reshape each to (N, 1, 10, 10), then stack → (N, 3, 10, 10) r1 = self.fc1.forward_reshaped(x_flat, self.grid_h, self.grid_w) r2 = self.fc2.forward_reshaped(x_flat, self.grid_h, self.grid_w) r3 = self.fc3.forward_reshaped(x_flat, self.grid_h, self.grid_w) stacked = torch.cat([r1, r2, r3], dim=1) # (N, 3, 10, 10) # Single Conv2D conv_out = self.relu(self.bn(self.conv(stacked))) # Global pool + classify #pooled = self.pool(conv_out).view(batch_size, -1) pooled = conv_out.view(batch_size, -1) logits = self.classifier(pooled) if return_intermediates: return logits, { 'bw1': bw1, 'bw2': bw2, 'bw3': bw3, 'stacked': stacked, 'conv_out': conv_out } return logits def train(model, train_loader, criterion, optimizer, epoch, device): """Single epoch training.""" model.train() running_loss = 0.0 correct = 0 total = 0 for batch_idx, (data, targets) in enumerate(train_loader): data, targets = data.to(device), targets.to(device) optimizer.zero_grad() outputs = model(data) loss = criterion(outputs, targets) loss.backward() optimizer.step() running_loss += loss.item() _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() if batch_idx % 100 == 0: acc = 100. * correct / total avg_loss = running_loss / (batch_idx + 1) print(f' Epoch {epoch} [{batch_idx}/{len(train_loader)}] ' f'Loss: {avg_loss:.4f} Acc: {acc:.2f}%') return running_loss / len(train_loader), 100. * correct / total def test(model, test_loader, criterion, device): """Evaluate model.""" model.eval() test_loss = 0 correct = 0 total = 0 with torch.no_grad(): for data, targets in test_loader: data, targets = data.to(device), targets.to(device) outputs = model(data) test_loss += criterion(outputs, targets).item() _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() test_loss /= len(test_loader) accuracy = 100. * correct / total return test_loss, accuracy def visualize_bandwidth(bandwidth, title="Bandwidth Distribution"): """Print bandwidth statistics.""" print(f"\n{title}:") print(f" Mean: {bandwidth.mean().item():.4f}") print(f" Std: {bandwidth.std().item():.4f}") print(f" Min: {bandwidth.min().item():.4f}") print(f" Max: {bandwidth.max().item():.4f}") def main(): # Configuration device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}\n") batch_size = 128 num_epochs = 30 learning_rate = 0.001 # Data transforms train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)) ]) test_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)) ]) # Load CIFAR-10 print("Loading CIFAR-10...") train_dataset = datasets.CIFAR10( root='../data', train=True, download=True, transform=train_transform ) test_dataset = datasets.CIFAR10( root='../data', train=False, download=True, transform=test_transform ) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) # Create model print("Creating BandwidthMLP model...") model = BandwidthMLP(img_size=32, num_classes=10).to(device) # Count parameters total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"Total parameters: {total_params:,}") print(f"Trainable parameters: {trainable_params:,}\n") # Loss and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs) # Training loop print("Starting training...\n" + "="*60) best_acc = 0.0 for epoch in range(1, num_epochs + 1): train_loss, train_acc = train(model, train_loader, criterion, optimizer, epoch, device) test_loss, test_acc = test(model, test_loader, criterion, device) scheduler.step() print(f"\nEpoch {epoch}/{num_epochs}") print(f" Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%") print(f" Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%") # Save best model if test_acc > best_acc: best_acc = test_acc torch.save(model.state_dict(), 'bandwidth_mlp_best.pt') print(f" *** New best model saved! (Acc: {best_acc:.2f}%)") print("="*60) # Final evaluation print(f"\n{'='*60}") print(f"Training complete!") print(f"Best test accuracy: {best_acc:.2f}%") # Load best model and show bandwidth stats print(f"\nLoading best model for bandwidth analysis...") model.load_state_dict(torch.load('bandwidth_mlp_best.pt', weights_only=True)) model.eval() # Show a sample prediction with torch.no_grad(): sample_data, sample_label = test_dataset[0] sample_input = sample_data.unsqueeze(0).to(device) logits, intermediates = model(sample_input, return_intermediates=True) pred = logits.argmax(1).item() classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'] print(f"\nSample prediction:") print(f" True label: {classes[sample_label]}") print(f" Predicted: {classes[pred]}") print(f"\nArchitecture: 3x Linear(3072→100) → reshape → (N,3,10,10) → Conv2d(3,64,k=3) → pool → classify") # Bandwidth statistics visualize_bandwidth(intermediates['bw1'], "FC1 Bandwidth (100 neurons)") visualize_bandwidth(intermediates['bw2'], "FC2 Bandwidth (100 neurons)") visualize_bandwidth(intermediates['bw3'], "FC3 Bandwidth (100 neurons)") if __name__ == "__main__": main()