File size: 8,137 Bytes
6eecae3 | 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 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 | """
test_robustness.py
Robustness evaluation for student models (Phase 2, Steps 2 & 3).
Tests trained models under various input corruptions and reports
localization error in meters (median, P90, P99).
Step 2 — AP Dropout:
- Remove each AP individually (zero out 1 of 4 channels)
- Random 1-AP dropout (random channel zeroed per sample)
- Random 2-AP dropout (2 random channels zeroed per sample)
Step 3 — Data Augmentation Robustness (optional):
- Gaussian noise (sigma=0.1, 0.2)
- AP attenuation (scale channel by 0.5)
- Phase perturbation (random rotation in complex-like space)
- Heatmap blur (Gaussian blur on input)
Usage:
python test_robustness.py
Edit CONFIG below to select model and weights path.
"""
import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as torchF
import numpy as np
from utils import localization_error
from data_loader import load_data
# ============================================================
# CONFIG
# ============================================================
MODEL_TYPE = 'mobilenet' # 'mobilenet', 'tinycnn', 'unet'
WEIGHTS_PATH = './weights/mobilenet_standalone_best.pth'
BATCH_SIZE = 32
# Data paths (Fig 10b — Jacobs Hall)
testpath = [
'./data/dataset_fov_test_jacobs_July28_2.mat',
'./data/dataset_non_fov_test_jacobs_July28_2.mat',
]
# ============================================================
# CORRUPTION FUNCTIONS
# ============================================================
def no_corruption(inputs):
"""Baseline: no corruption."""
return inputs
def drop_ap(inputs, ap_index):
"""Zero out a specific AP channel."""
corrupted = inputs.clone()
corrupted[:, ap_index, :, :] = 0.0
return corrupted
def random_drop_1ap(inputs):
"""Randomly zero 1 AP per sample."""
corrupted = inputs.clone()
for i in range(corrupted.shape[0]):
ap = torch.randint(0, 4, (1,)).item()
corrupted[i, ap, :, :] = 0.0
return corrupted
def random_drop_2ap(inputs):
"""Randomly zero 2 APs per sample."""
corrupted = inputs.clone()
for i in range(corrupted.shape[0]):
aps = torch.randperm(4)[:2]
for ap in aps:
corrupted[i, ap, :, :] = 0.0
return corrupted
def gaussian_noise(inputs, sigma=0.1):
"""Add Gaussian noise."""
return inputs + torch.randn_like(inputs) * sigma
def ap_attenuation(inputs, scale=0.5):
"""Scale all channels by a factor (simulates weaker signal)."""
return inputs * scale
def single_ap_attenuation(inputs, ap_index, scale=0.3):
"""Attenuate a single AP channel."""
corrupted = inputs.clone()
corrupted[:, ap_index, :, :] *= scale
return corrupted
def gaussian_blur(inputs, kernel_size=5, sigma=1.0):
"""Apply Gaussian blur to each channel."""
# Create 1D Gaussian kernel
x = torch.arange(kernel_size, dtype=torch.float32, device=inputs.device) - kernel_size // 2
kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2)
kernel_1d = kernel_1d / kernel_1d.sum()
# Make 2D kernel
kernel_2d = kernel_1d[:, None] * kernel_1d[None, :]
kernel_2d = kernel_2d.expand(inputs.shape[1], 1, kernel_size, kernel_size)
pad = kernel_size // 2
return torchF.conv2d(inputs, kernel_2d, padding=pad, groups=inputs.shape[1])
# ============================================================
# MAIN
# ============================================================
def load_model(model_type, weights_path, device):
if model_type == 'mobilenet':
from student_mobilenet import MobileNetStudent
model = MobileNetStudent(input_nc=4)
elif model_type == 'tinycnn':
from student_tinycnn import TinyCNNStudent
model = TinyCNNStudent(input_nc=4)
elif model_type == 'unet':
from student_unet import MobileNetUNetStudent
model = MobileNetUNetStudent(input_nc=4)
else:
raise ValueError(f'Unknown model type: {model_type}')
model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
model.to(device)
model.eval()
return model
def evaluate_condition(model, test_loader, corrupt_fn, device):
"""Run model on corrupted inputs, return error stats in meters."""
all_errors = []
with torch.no_grad():
for data in test_loader:
inputs = data[0].to(device)
labels = data[1].to(device)
corrupted = corrupt_fn(inputs)
output = model(corrupted)
errors = localization_error(
output.cpu().numpy(),
labels.cpu().numpy(),
scale=0.1,
)
all_errors.extend(errors)
all_errors = np.array(all_errors)
return {
'median': np.median(all_errors),
'p90': np.percentile(all_errors, 90),
'p99': np.percentile(all_errors, 99),
'mean': np.mean(all_errors),
}
def main():
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f'Device: {device}')
print(f'Model: {MODEL_TYPE}, Weights: {WEIGHTS_PATH}')
# Load model
model = load_model(MODEL_TYPE, WEIGHTS_PATH, device)
n_params = sum(p.numel() for p in model.parameters())
print(f'Parameters: {n_params:,}')
# Load test data
print('Loading test data...')
test_w_list, test_lbl_list = [], []
for fp in testpath:
_, w, lbl = load_data(fp)
if lbl.dim() == 3:
lbl = lbl.unsqueeze(1)
test_w_list.append(w)
test_lbl_list.append(lbl)
test_w = torch.cat(test_w_list, dim=0)
test_lbl = torch.cat(test_lbl_list, dim=0)
print(f'Test samples: {test_w.shape[0]}')
test_data = torch.utils.data.TensorDataset(test_w, test_lbl)
test_loader = torch.utils.data.DataLoader(
test_data, batch_size=BATCH_SIZE, shuffle=False,
num_workers=4, pin_memory=True,
)
# Define all test conditions
conditions = {
# Step 2: AP Dropout
'Normal (no corruption)': no_corruption,
'Drop AP 0': lambda x: drop_ap(x, 0),
'Drop AP 1': lambda x: drop_ap(x, 1),
'Drop AP 2': lambda x: drop_ap(x, 2),
'Drop AP 3': lambda x: drop_ap(x, 3),
'Random drop 1 AP': random_drop_1ap,
'Random drop 2 APs': random_drop_2ap,
# Step 3: Augmentation robustness
'Gaussian noise (sigma=0.1)': lambda x: gaussian_noise(x, sigma=0.1),
'Gaussian noise (sigma=0.2)': lambda x: gaussian_noise(x, sigma=0.2),
'Attenuation (0.5x all)': lambda x: ap_attenuation(x, scale=0.5),
'Attenuate AP 0 (0.3x)': lambda x: single_ap_attenuation(x, 0, 0.3),
'Gaussian blur (k=5, s=1)': lambda x: gaussian_blur(x, kernel_size=5, sigma=1.0),
}
# Run all conditions
print(f'\n{"="*75}')
print(f' ROBUSTNESS EVALUATION — {MODEL_TYPE.upper()}')
print(f'{"="*75}')
print(f'{"Condition":<30} {"Median":>8} {"P90":>8} {"P99":>8} {"Mean":>8}')
print(f'{"-"*75}')
results = {}
for name, corrupt_fn in conditions.items():
stats = evaluate_condition(model, test_loader, corrupt_fn, device)
results[name] = stats
print(f'{name:<30} {stats["median"]:>7.3f}m {stats["p90"]:>7.3f}m {stats["p99"]:>7.3f}m {stats["mean"]:>7.3f}m')
print(f'{"="*75}')
# Save results
save_path = f'robustness_{MODEL_TYPE}.txt'
with open(save_path, 'w') as f:
f.write(f'Model: {MODEL_TYPE}\n')
f.write(f'Weights: {WEIGHTS_PATH}\n')
f.write(f'Test samples: {test_w.shape[0]}\n\n')
f.write(f'{"Condition":<30} {"Median":>8} {"P90":>8} {"P99":>8} {"Mean":>8}\n')
f.write(f'{"-"*62}\n')
for name, stats in results.items():
f.write(f'{name:<30} {stats["median"]:>7.3f}m {stats["p90"]:>7.3f}m {stats["p99"]:>7.3f}m {stats["mean"]:>7.3f}m\n')
print(f'\nResults saved to {save_path}')
if __name__ == '__main__':
main()
|