3 datasets, 4 models, 4 hardening techniques

Author: jamilgafur

manuscripts/Poison26/bin/train/AUDIOMNIST/MobileNet/MobileNet.py

@@ -9,14 +9,12 @@ class DepthwiseSeparableConv(nn.Module):
     def __init__(self, in_channels, out_channels, stride=1):
         super(DepthwiseSeparableConv, self).__init__()
 
-        # Depthwise layer with BN and ReLU6
         self.depthwise = nn.Sequential(
             nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False),
             nn.BatchNorm2d(in_channels),
             nn.ReLU6(inplace=True)
         )
 
-        # Pointwise layer with BN and ReLU6
         self.pointwise = nn.Sequential(
             nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
             nn.BatchNorm2d(out_channels),
@@ -32,10 +30,9 @@ def forward(self, x):
 # MobileNetV1 (Model Definition)
 # -------------------------
 class MobileNet(nn.Module):
-    def __init__(self, one_batch=None, num_classes=1000):
+    def __init__(self, one_batch=None, num_classes=10):
         super(MobileNet, self).__init__()
 
-        # Handle dynamic input sizes
         if one_batch is not None:
             _, in_channels, H, W = one_batch.shape
             self.input_channels = in_channels
@@ -44,18 +41,12 @@ def __init__(self, one_batch=None, num_classes=1000):
             self.input_channels = 3
             self.input_size = (3, 224, 224)
 
-        # -------------------------
-        # Stem
-        # -------------------------
         self.stem = nn.Sequential(OrderedDict([
             ('conv1', nn.Conv2d(self.input_channels, 32, kernel_size=3, stride=2, padding=1, bias=False)),
             ('bn1', nn.BatchNorm2d(32)),
             ('relu1', nn.ReLU6(inplace=True)),
         ]))
 
-        # -------------------------
-        # Full MobileNetV1 Architecture
-        # -------------------------
         layers = [
             DepthwiseSeparableConv(32, 64, stride=1),
             DepthwiseSeparableConv(64, 128, stride=2),
@@ -65,29 +56,19 @@ def __init__(self, one_batch=None, num_classes=1000):
             DepthwiseSeparableConv(256, 512, stride=2)
         ]
 
-        # 5x repeating blocks of 512 channels
         for _ in range(5):
             layers.append(DepthwiseSeparableConv(512, 512, stride=1))
 
-        # Final expansion to 1024 channels
         layers.extend([
             DepthwiseSeparableConv(512, 1024, stride=2),
             DepthwiseSeparableConv(1024, 1024, stride=1)
         ])
 
-        # Pack the layers into an nn.Sequential for cleaner forward pass
         self.features = nn.Sequential(*layers)
-
-        # -------------------------
-        # Classifier Setup
-        # -------------------------
         self.pool = nn.AdaptiveAvgPool2d(1)
         self.fc_input_features = self._get_flattened_feature_size(one_batch)
         self.fc = nn.Linear(self.fc_input_features, num_classes)
 
-    # -------------------------
-    # Compute FC feature size dynamically
-    # -------------------------
     def _get_flattened_feature_size(self, one_batch):
         was_training = self.training
         self.eval()
@@ -109,13 +90,10 @@ def _get_flattened_feature_size(self, one_batch):
 
         return out_features
 
-    # -------------------------
-    # Forward
-    # -------------------------
     def forward(self, x):
         x = self.stem(x)
         x = self.features(x)
         x = self.pool(x)
         x = torch.flatten(x, 1)
         x = self.fc(x)
-        return x
+        return x

manuscripts/Poison26/bin/train/AUDIOMNIST/MobileNet/model_aug.py

@@ -7,51 +7,46 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import torch.optim as optim
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, Subset
 
 import librosa
 from sklearn.metrics import roc_auc_score, average_precision_score
 
 import hashlib
 import csv
-
-# ----------------------------
-# Import the 2D Vision Models
-# ----------------------------
-from RegNetX import RegNetX_400MF
-from MobileNet import MobileNet
-# from ConvNetX import ConvNeXt
+import random
 
 # ----------------------------
 # Constants
 # ----------------------------
 SAMPLING_RATE = 16000
 NUM_CLASSES = 10
 MAX_AUDIO_LENGTH = 16000
-BATCH_SIZE = 32
 
 # ----------------------------
 # Audio Preprocessing
 # ----------------------------
 def normalize_audio(x):
-    return x / np.max(np.abs(x))
+    max_val = np.max(np.abs(x))
+    return x / max_val if max_val > 0 else x
 
 def pad_audio(audio, max_len=MAX_AUDIO_LENGTH):
     return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant')
 
 # ----------------------------
-# Dataset
+# Dataset & Wrapper
 # ----------------------------
-class AudioMNISTDataset(Dataset):
+class AudioMNISTBaseDataset(Dataset):
+    """Loads all audio into memory ONCE. Returns raw numpy arrays."""
     def __init__(self, data_path):
         self.data = []
         self.labels = []
 
         wav_files = glob.glob(os.path.join(data_path, '*', '*.wav'))
-        # Deterministic shuffle using md5 hash of path
         wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest())
-        self.wav_files = wav_files.copy()  # store for TSV
+        self.wav_files = wav_files.copy()
 
         for audio_path in tqdm(wav_files, desc="Loading audio files"):
             audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE)
@@ -65,12 +60,60 @@ def __len__(self):
         return len(self.data)
 
     def __getitem__(self, idx):
-        audio = torch.tensor(self.data[idx], dtype=torch.float32).unsqueeze(0)
-        label = self.labels[idx]
-        return audio, label
+        return self.data[idx], self.labels[idx]
+
+class AudioSubsetWrapper(Dataset):
+    """Wraps a subset to apply dynamic augmentation and convert to Tensors."""
+    def __init__(self, subset, augment=False):
+        self.subset = subset
+        self.augment = augment
+
+    def __len__(self):
+        return len(self.subset)
+
+    def apply_augmentation(self, x):
+        if random.random() < 0.5:
+            x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) # noise
+        if random.random() < 0.5:
+            x = np.roll(x, np.random.randint(-200, 200)) # time shift
+        if random.random() < 0.5:
+            x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) # random gain
+        return x
+
+    def __getitem__(self, idx):
+        x, y = self.subset[idx]
+        if self.augment:
+            x = self.apply_augmentation(x)
+        x = torch.tensor(x, dtype=torch.float32).unsqueeze(0)  # (1, length)
+        return x, y
+
+def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"):
+    base_dataset = AudioMNISTBaseDataset(data_path)
+    
+    train_size = int(0.8 * len(base_dataset))
+    train_indices = list(range(0, train_size))
+    test_indices  = list(range(train_size, len(base_dataset)))
+
+    # Isolate augmentation using the wrapper
+    train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train)
+    test_dataset  = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    test_loader  = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+    with open(split_tsv, "w", newline="") as f:
+        writer = csv.writer(f, delimiter="\t")
+        writer.writerow(["index", "split", "label", "file_path"])
+        for idx in train_indices:
+            writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]])
+        for idx in test_indices:
+            writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]])
+    print(f"Saved split information to {split_tsv}")
+
+    return train_loader, test_loader
 
 # ----------------------------
-# Model Definition (Updated Wrapper)
+# Model Definition
 # ----------------------------
 class AudioMNISTModel(nn.Module):
     """
@@ -97,41 +140,12 @@ def forward(self, x):
         x = x.view(x.size(0), *self.reshape_dims)
         return self.backbone(x)
 
-# ----------------------------
-# Load Data
-# ----------------------------
-def load_data(data_path, batch_size, split_tsv="split_indices_model1.tsv"):
-    dataset = AudioMNISTDataset(data_path)
-    # Fixed 80/20 split (after deterministic shuffle)
-    train_size = int(0.8 * len(dataset))
-    train_indices = list(range(0, train_size))
-    test_indices  = list(range(train_size, len(dataset)))
-
-    train_dataset = torch.utils.data.Subset(dataset, range(0, train_size))
-    test_dataset  = torch.utils.data.Subset(dataset, range(train_size, len(dataset)))
-
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-    test_loader  = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
-
-    # --- Write split info to TSV ---
-    with open(split_tsv, "w", newline="") as f:
-        writer = csv.writer(f, delimiter="\t")
-        writer.writerow(["index", "split", "label", "file_path"])
-        for idx in train_indices:
-            writer.writerow([idx, "train", dataset.labels[idx], dataset.wav_files[idx]])
-        for idx in test_indices:
-            writer.writerow([idx, "test", dataset.labels[idx], dataset.wav_files[idx]])
-    print(f"Saved split information to {split_tsv}")
-    
-    return train_loader, test_loader
-
 # ----------------------------
 # Training loop
 # ----------------------------
 def train(model, train_loader, device, epochs=10, lr=0.001):
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=lr)
-
     model.to(device)
     model.train()
 
@@ -144,7 +158,7 @@ def train(model, train_loader, device, epochs=10, lr=0.001):
         for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"):
             images, labels = images.to(device), labels.to(device)
 
-            optimizer.zero_grad()
+            optimizer.zero_grad(set_to_none=True)
             outputs = model(images)
             loss = criterion(outputs, labels)
             loss.backward()
@@ -160,7 +174,6 @@ def train(model, train_loader, device, epochs=10, lr=0.001):
         elapsed = time.time() - start_time
         print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}")
 
-
 # ----------------------------
 # Evaluation
 # ----------------------------
@@ -196,9 +209,8 @@ def evaluate_model(model, test_loader, device):
 
     y_true = np.array(y_true)
     y_pred = np.array(y_pred)
-
-    # compute AUROC and AUPRC
-    y_true_onehot = np.eye(10)[y_true]
+    y_true_onehot = np.eye(NUM_CLASSES)[y_true]
+    
     auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr")
     auprc = average_precision_score(y_true_onehot, y_pred)
 
@@ -207,42 +219,29 @@ def evaluate_model(model, test_loader, device):
     print(f"Test auROC: {auroc:.4f}")
     print(f"Test auPRC: {auprc:.4f}")
 
-
 # ----------------------------
 # Main
 # ----------------------------
 def main():
-    parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with Augmentation")
-    parser.add_argument("--output", type=str, default="mnist_model_aug.pt", help="Model output name")
+    parser = argparse.ArgumentParser(description="AudioMNIST Augmented Training")
+    parser.add_argument("--data", type=str, default="./data/AudioMNIST", help="Path to dataset")
+    parser.add_argument("--output", type=str, default="audiomnist_aug.pt", help="Model output name")
     parser.add_argument("--batch-size", type=int, default=64)
-    parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs")
+    parser.add_argument("--epochs", type=int, default=10)
     args = parser.parse_args()
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = AudioMNISTModel(backbone_class=MobileNet, num_classes=NUM_CLASSES)
 
-    
-    # Load data
-    train_loader, test_loader = load_data(batch_size=args.batch_size)
-    
-    # Initialize model
-    # Pass a dummy batch to configure the MobileNet stem for 1-channel MNIST images
-    # and properly calculate the fully-connected layer inputs for 28x28 resolution.
-    
-    dummy_batch = train_loader.dataset[0][0].unsqueeze(0) 
-    model = MobileNet(one_batch=dummy_batch, num_classes=10)
+    # ENABLE DATA AUGMENTATION HERE
+    train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=True, split_tsv="split_indices_aug.tsv")
 
-
-    # Train
     train(model, train_loader, device, epochs=args.epochs)
-
-    # Save model
     torch.save(model.state_dict(), args.output)
     print(f"Model saved to {args.output}")
 
-    # Evaluate
-    print("Model statistics on test dataset")
+    print("Model statistics on clean test dataset")
     evaluate_model(model, test_loader, device)
 
-
 if __name__ == "__main__":
     main()