Spaces:
Sleeping
Sleeping
File size: 5,032 Bytes
95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 95382f9 da4f171 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 |
"""
Attack CNN Model Architectures
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class StandardCNN(nn.Module):
"""
Standard CNN Model (Original)
Architecture: 3 Conv blocks with BatchNorm + 3 FC layers
Parameters: ~817K
"""
def __init__(self, num_classes=10, dropout_rate=0.5):
super(StandardCNN, self).__init__()
# First convolutional block
self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)
self.bn1 = nn.BatchNorm2d(32)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
# Second convolutional block
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
self.bn2 = nn.BatchNorm2d(64)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
# Third convolutional block
self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
self.bn3 = nn.BatchNorm2d(128)
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
# Calculate the flattened size after convolutions
self.flattened_size = 128 * 3 * 3 # 28x28 -> 14x14 -> 7x7 -> 3x3
# Fully connected layers
self.fc1 = nn.Linear(self.flattened_size, 512)
self.dropout1 = nn.Dropout(dropout_rate)
self.fc2 = nn.Linear(512, 256)
self.dropout2 = nn.Dropout(dropout_rate)
self.fc3 = nn.Linear(256, num_classes)
def forward(self, x, return_logits=False):
# Conv block 1
x = self.conv1(x)
x = self.bn1(x)
x = F.relu(x)
x = self.pool1(x)
# Conv block 2
x = self.conv2(x)
x = self.bn2(x)
x = F.relu(x)
x = self.pool2(x)
# Conv block 3
x = self.conv3(x)
x = self.bn3(x)
x = F.relu(x)
x = self.pool3(x)
# Flatten
x = x.view(x.size(0), -1)
# FC layers
x = F.relu(self.fc1(x))
x = self.dropout1(x)
x = F.relu(self.fc2(x))
x = self.dropout2(x)
logits = self.fc3(x)
if return_logits:
return logits
return F.softmax(logits, dim=1)
class LighterCNN(nn.Module):
"""
Lighter CNN Model
Architecture: 3 Conv blocks with fewer filters + Global Average Pooling
Parameters: ~94K
"""
def __init__(self, num_classes=10, dropout_rate=0.5):
super(LighterCNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
self.bn1 = nn.BatchNorm2d(32)
self.pool1 = nn.MaxPool2d(2,2)
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.bn2 = nn.BatchNorm2d(64)
self.pool2 = nn.MaxPool2d(2,2)
self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
self.bn3 = nn.BatchNorm2d(128)
self.pool3 = nn.MaxPool2d(2,2) # 28->14->7->3
self.gap = nn.AdaptiveAvgPool2d(1) # (B,128,1,1)
self.fc = nn.Linear(128, num_classes)
def forward(self, x, return_logits=False):
x = self.pool1(F.relu(self.bn1(self.conv1(x))))
x = self.pool2(F.relu(self.bn2(self.conv2(x))))
x = self.pool3(F.relu(self.bn3(self.conv3(x))))
x = self.gap(x).view(x.size(0), -1) # (B,128)
logits = self.fc(x)
return logits if return_logits else F.softmax(logits, dim=1)
class DepthwiseSeparableConv(nn.Module):
def __init__(self, in_ch, out_ch, stride=1):
super(DepthwiseSeparableConv, self).__init__()
self.dw = nn.Conv2d(in_ch, in_ch, 3, stride=stride, padding=1,
groups=in_ch, bias=False) # depthwise
self.pw = nn.Conv2d(in_ch, out_ch, 1, bias=False) # pointwise
self.bn = nn.BatchNorm2d(out_ch)
def forward(self, x):
x = self.dw(x)
x = self.pw(x)
return F.relu(self.bn(x), inplace=True)
class DepthwiseCNN(nn.Module):
"""
Depthwise Separable CNN
Ultra-efficient using Depthwise Separable Convolutions
Parameters: ~1.4K
"""
def __init__(self, num_classes=10, dropout_rate=0.5):
super(DepthwiseCNN, self).__init__()
# Stem: 1 -> 8, reduce size with stride=2 (28->14)
self.stem = nn.Sequential(
nn.Conv2d(1, 8, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(8),
nn.ReLU(inplace=True),
)
# DS blocks
self.ds1 = DepthwiseSeparableConv(8, 16, stride=1)
self.ds2 = DepthwiseSeparableConv(16, 32, stride=2) # 14->7
self.gap = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Linear(32, num_classes)
def forward(self, x, return_logits=False):
x = self.stem(x) # B, 8, 14, 14
x = self.ds1(x) # B,16,14,14
x = self.ds2(x) # B,32, 7, 7
x = self.gap(x).flatten(1) # B,32
logits = self.fc(x) # B,10
return logits if return_logits else F.softmax(logits, dim=1)
|