O1: wrn2 fixes

one_run_audit/audit.py

@@ -15,12 +15,15 @@ from torchvision.datasets import CIFAR10
 import pytorch_lightning as pl
 import opacus
 import random
+from tqdm import tqdm
 from opacus.validators import ModuleValidator
 from opacus.utils.batch_memory_manager import BatchMemoryManager
 from WideResNet import WideResNet
 from equations import get_eps_audit
 import student_model
 import fast_model
+import convnet_classifier
+import wrn
 import warnings
 warnings.filterwarnings("ignore")
 
@@ -230,8 +233,10 @@ def get_dataloaders_raw(m=1000, train_batch_size=512, test_batch_size=10):
 
     train_x = preprocess_data(train_x)
     test_x = preprocess_data(test_x)
+    attack_x = preprocess_data(attack_x)
     train_y = torch.tensor(train_y)
     test_y = torch.tensor(test_y)
+    attack_y = torch.tensor(attack_y)
 
     train_dl = DataLoader(
         TensorDataset(train_x, train_y.long()),
@@ -246,7 +251,7 @@ def get_dataloaders_raw(m=1000, train_batch_size=512, test_batch_size=10):
         shuffle=True,
         num_workers=4
     )
-    return train_dl, test_dl, train_x
+    return train_dl, test_dl, train_x, attack_x.numpy(), attack_y.numpy(), S
 
 def evaluate_on(model, dataloader):
     correct = 0
@@ -398,6 +403,70 @@ def load(hp, model_path, train_dl):
     return model_init, model, adv_points, adv_labels, S
 
 
+def train_wrn2(hp, train_dl, test_dl):
+    model = wrn.WideResNet(16, 10, 4)
+    model = model.to(DEVICE)
+    #model = ModuleValidator.fix(model)
+    ModuleValidator.validate(model, strict=True)
+    model_init = copy.deepcopy(model)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(
+        model.parameters(),
+        lr=0.12,
+        momentum=0.9,
+        weight_decay=1e-4
+    )
+    scheduler = MultiStepLR(
+        optimizer,
+        milestones=[int(i * hp['epochs']) for i in [0.3, 0.6, 0.8]],
+        gamma=0.1
+    )
+
+    print(f"Training with {hp['epochs']} epochs")
+
+    if hp['epsilon'] is not None:
+        privacy_engine = opacus.PrivacyEngine()
+        model, optimizer, train_loader = privacy_engine.make_private_with_epsilon(
+            module=model,
+            optimizer=optimizer,
+            data_loader=train_dl,
+            epochs=hp['epochs'],
+            target_epsilon=hp['epsilon'],
+            target_delta=hp['delta'],
+            max_grad_norm=hp['norm'],
+        )
+
+        print(f"DP epsilon = {hp['epsilon']}, delta = {hp['delta']}")
+        print(f"Using sigma={optimizer.noise_multiplier} and C = norm = {hp['norm']}")
+
+        with BatchMemoryManager(
+            data_loader=train_loader,
+            max_physical_batch_size=10,  # 1000 ~= 9.4GB vram
+            optimizer=optimizer
+        ) as memory_safe_data_loader:
+            best_test_set_accuracy = train_no_cap(
+                model,
+                hp,
+                memory_safe_data_loader,
+                test_dl,
+                optimizer,
+                criterion,
+                scheduler,
+            )
+    else:
+        print("Training without differential privacy")
+        best_test_set_accuracy = train_no_cap(
+            model,
+            hp,
+            train_dl,
+            test_dl,
+            optimizer,
+            criterion,
+            scheduler,
+        )
+
+    return model_init, model
 
 def train_small(hp, train_dl, test_dl):
     model = student_model.Model(num_classes=10).to(DEVICE)
@@ -460,7 +529,7 @@ def train_small(hp, train_dl, test_dl):
 
     return model_init, model
 
-def train_fast(hp):
+def train_fast(hp, train_dl, test_dl, train_x):
     epochs = hp['epochs']
     momentum = 0.9
     weight_decay = 0.256
@@ -472,8 +541,6 @@ def train_fast(hp):
     print("=========================")
     print("Training a fast model")
     print("=========================")
-    train_dl, test_dl, train_x = get_dataloaders_raw(hp['target_points'])
-
     weights = fast_model.patch_whitening(train_x[:10000, :, 4:-4, 4:-4])
     model = fast_model.Model(weights, c_in=3, c_out=10, scale_out=0.125)
 
@@ -604,12 +671,12 @@ def main():
     parser.add_argument('--cuda', type=int, help='gpu index', required=False)
     parser.add_argument('--epsilon', type=float, help='dp epsilon', required=False, default=None)
     parser.add_argument('--m', type=int, help='number of target points', required=True)
-    parser.add_argument('--k', type=int, help='number of symmetric guesses', required=True)
     parser.add_argument('--epochs', type=int, help='number of epochs', required=True)
     parser.add_argument('--load', type=Path, help='number of epochs', required=False)
     parser.add_argument('--studentraw', action='store_true', help='train a raw student', required=False)
     parser.add_argument('--distill', action='store_true', help='train a raw student', required=False)
-    parser.add_argument('--fast', action='store_true', help='train a the fast model', required=False)
+    parser.add_argument('--fast', action='store_true', help='train the fast model', required=False)
+    parser.add_argument('--wrn2', action='store_true', help='Train a groupnormed wrn', required=False)
     args = parser.parse_args()
 
     if torch.cuda.is_available() and args.cuda:
@@ -631,8 +698,6 @@ def main():
         "norm": args.norm,
         "batch_size": 4096,
         "epochs": args.epochs,
-        "k+": args.k,
-        "k-": args.k,
         "p_value": 0.05,
     }
 
@@ -652,12 +717,16 @@ def main():
         model_init, model_trained, adv_points, adv_labels, S = load(hp, args.load, train_dl)
         test_dl = None
     elif args.fast:
-        train_dl, test_dl, _ = get_dataloaders_raw(hp['target_points'])
+        train_dl, test_dl, train_x, adv_points, adv_labels, S = get_dataloaders_raw(hp['target_points'])
-        model_init, model_trained = train_fast(hp)
+        model_init, model_trained = train_fast(hp, train_dl, test_dl, train_x)
-        exit(1)
     else:
         train_dl, test_dl, pure_train_dl, adv_points, adv_labels, S = get_dataloaders3(hp['target_points'], hp['batch_size'])
-        if args.studentraw:
+        if args.wrn2:
+            print("=========================")
+            print("Training wrn2 model from meta")
+            print("=========================")
+            model_init, model_trained = train_wrn2(hp, train_dl, test_dl)
+        elif args.studentraw:
             print("=========================")
             print("Training a raw student model")
             print("=========================")
@@ -711,13 +780,18 @@ def main():
 
             scores.append(((init_loss - trained_loss).item(), is_in))
 
+    print(f"Top 10 unsorted scores: {scores[:10]}")
+    print(f"Btm 10 unsorted scores: {scores[-10:]}")
     scores = sorted(scores, key=lambda x: x[0])
+    print(f"Top 10 sorted scores: {scores[:10]}")
+    print(f"Btm 10 sorted scores: {scores[-10:]}")
     scores = np.array([x[1] for x in scores])
 
-    print(scores[:10])
-
     audits = (0, 0, 0, 0)
-    for k in [10, 20, 50, 100, 200, 300, 500, 800, 1000, 1200, 1400, 1600, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500]:
+    k_schedule = np.linspace(1, hp['target_points']//2, 40)
+    k_schedule = np.floor(k_schedule).astype(int)
+
+    for k in tqdm(k_schedule):
         correct = np.sum(~scores[:k]) + np.sum(scores[-k:])
         total = len(scores)