Update train_robust.py

train_robust.py

@@ -15,7 +15,7 @@ SAVE_EVERY = 10
 CHECKPOINT_DIR = "./checkpoints"
 BEST_MODEL_DIR = "./nula-best-model"
 
-def train_one_epoch(model, loader, optimizer, device, grad_clip=1.0):
+def train_one_epoch(model, loader, optimizer, device, mean, std, blur, grad_clip=1.0):
     model.train()
     
     total_loss = 0.0
@@ -34,7 +34,7 @@ def train_one_epoch(model, loader, optimizer, device, grad_clip=1.0):
 
         if mask_aug.any():
             with pt.no_grad():
-                x_image = x * STD + MEAN
+                x_image = x * std + mean
                 choices = pt.randint(0, 3, (B,), device=x.device)
                 mask_resize = mask_aug & (choices == 0)
                 if mask_resize.any():
@@ -58,11 +58,11 @@ def train_one_epoch(model, loader, optimizer, device, grad_clip=1.0):
                 mask_blur = mask_aug & (choices == 2)
                 if mask_blur.any():
                     x_subset = x_image[mask_blur]
-                    x_down = GLOBAL_POOL_BLUR(x_subset)
+                    x_down = blur(x_subset)
                     x_up = F.interpolate(x_down, size=x_subset.shape[-2:], mode="bilinear", align_corners=False)
                     x_image[mask_blur] = x_up
 
-                x = (x_image - MEAN) / STD
+                x = (x_image - mean) / std
 
         out = model(pixel_values=x, labels=y)
         loss = out.loss
@@ -110,16 +110,24 @@ if __name__ == "__main__":
     cosine = CosineAnnealingLR(optimizer, T_max=45)
     scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[5])
     
-    MEAN             = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
-    STD              = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
+    MEAN = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
+    STD = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
     GLOBAL_POOL_BLUR = BlurPool2d(channels=cfg.in_channels, stride=2).to(DEVICE)
     best_val_acc = 0.0
 
     os.makedirs(CHECKPOINT_DIR, exist_ok=True)
 
     for epoch in range(1, NUM_EPOCHS + 1):
-        train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, DEVICE)
-        val_loss, val_acc     = evaluate(model, test_loader, DEVICE)
+        train_loss, train_acc = train_one_epoch(
+            model, 
+            train_loader, 
+            optimizer, 
+            DEVICE, 
+            MEAN, 
+            STD, 
+            GLOBAL_POOL_BLUR
+        )
+        val_loss, val_acc = evaluate(model, test_loader, DEVICE)
         scheduler.step()
 
         if epoch % SAVE_EVERY == 0: