Merge pull request #4 from YmShan/main

Directly runnable MAE pre training.

Author: bollossom

SDT_V3/Classification/Model_Large/MAE_SDT.py

@@ -280,7 +280,7 @@ def __init__(
             self.layer_scale2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
 
     def forward(self, x):
-        # T, B, C, N = x.shape
+        T, B, C, N = x.shape
         if self.model=="base":
             x= x + self.rep_conv(self.lif(x).flatten(0, 1)).reshape(T, B, C, N)
         # TODO: need channel-wise layer scale, init as 1e-6

SDT_V3/Classification/Model_Large/engine_pretrain.py

@@ -55,7 +55,7 @@ def train_one_epoch(model: torch.nn.Module,
             optimizer.zero_grad()
 
         torch.cuda.synchronize()
-        functional.reset_net(model)
+        # functional.reset_net(model)
         metric_logger.update(loss=loss_value)
 
         lr = optimizer.param_groups[0]["lr"]
@@ -74,4 +74,4 @@ def train_one_epoch(model: torch.nn.Module,
     # gather the stats from all processes
     metric_logger.synchronize_between_processes()
     print("Averaged stats:", metric_logger)
-    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}