feat: speed up training 20260828

diffusion_planner/diffusion_planner/config/model_config.py

@@ -47,7 +47,12 @@ class ModelConfig:
     # ---------------------------------------------------------
     # Model Architecture
     # ---------------------------------------------------------
-    encoder_mixer_depth: int = 6
+    # Per-element MLP-Mixer size. These blocks run on every element (~560 per sample)
+    # at every point (20-31), so their activations dominate encoder time and memory:
+    # cost scales as depth * mixer dims, not as hidden_dim. 128/64/6 is what ran before.
+    encoder_mixer_hidden_dim: int = 32
+    encoder_mixer_token_dim: int = 32
+    encoder_mixer_depth: int = 1
     encoder_fusion_depth: int = 6
     decoder_depth: int = 3
     num_heads: int = 8

diffusion_planner/diffusion_planner/config/train_config.py

@@ -35,7 +35,7 @@ class TrainConfig(ClosedLoopConfig, ScenarioOpenLoopConfig, ModelConfig):
     # ---------------------------------------------------------
     # DataLoader Parameters
     # ---------------------------------------------------------
-    batch_size: int = cli("batch size across all GPUs", default=512)
+    batch_size: int = cli("batch size across all GPUs", default=4096)
     num_workers: int = 8
     pin_mem: bool = True
 
@@ -91,6 +91,7 @@ class TrainConfig(ClosedLoopConfig, ScenarioOpenLoopConfig, ModelConfig):
     device: str = "cuda"
     use_ema: bool = True
     ema_decay: float = 0.999
+    use_amp: bool = cli("train with Automatic Mixed Precision (bf16 autocast)", default=True)
     resume_model_path: Optional[str] = cli(
         "resume training from this .pth", default=None, path=True
     )

diffusion_planner/diffusion_planner/model/module/encoder.py

@@ -70,12 +70,16 @@ def __init__(self, config):
             drop_path_rate=config.encoder_drop_path_rate,
             hidden_dim=config.hidden_dim,
             depth=config.encoder_mixer_depth,
+            mixer_hidden_dim=config.encoder_mixer_hidden_dim,
+            mixer_token_dim=config.encoder_mixer_token_dim,
         )
         self.neighbor_encoder = NeighborEncoder(
             config.time_len,
             drop_path_rate=config.encoder_drop_path_rate,
             hidden_dim=config.hidden_dim,
             depth=config.encoder_mixer_depth,
+            mixer_hidden_dim=config.encoder_mixer_hidden_dim,
+            mixer_token_dim=config.encoder_mixer_token_dim,
         )
         self.static_encoder = StaticEncoder(
             config.static_objects_state_dim,
@@ -88,20 +92,26 @@ def __init__(self, config):
             drop_path_rate=config.encoder_drop_path_rate,
             hidden_dim=config.hidden_dim,
             depth=config.encoder_mixer_depth,
+            mixer_hidden_dim=config.encoder_mixer_hidden_dim,
+            mixer_token_dim=config.encoder_mixer_token_dim,
         )
         self.route_encoder = LaneEncoder(
             config.route_len,
             class_type=CLASS_TYPE_ROUTE,
             drop_path_rate=config.encoder_drop_path_rate,
             hidden_dim=config.hidden_dim,
             depth=config.encoder_mixer_depth,
+            mixer_hidden_dim=config.encoder_mixer_hidden_dim,
+            mixer_token_dim=config.encoder_mixer_token_dim,
         )
         self.polygon_encoder = LineEncoder(
             config.polygon_len,
             class_type=CLASS_TYPE_POLYGON,
             drop_path_rate=config.encoder_drop_path_rate,
             hidden_dim=config.hidden_dim,
             depth=config.encoder_mixer_depth,
+            mixer_hidden_dim=config.encoder_mixer_hidden_dim,
+            mixer_token_dim=config.encoder_mixer_token_dim,
             point_dim=2 + POLYGON_TYPE_NUM,
         )
         self.line_string_encoder = LineEncoder(
@@ -110,6 +120,8 @@ def __init__(self, config):
             drop_path_rate=config.encoder_drop_path_rate,
             hidden_dim=config.hidden_dim,
             depth=config.encoder_mixer_depth,
+            mixer_hidden_dim=config.encoder_mixer_hidden_dim,
+            mixer_token_dim=config.encoder_mixer_token_dim,
             point_dim=2 + LINE_STRING_TYPE_NUM,
         )
         self.goal_pose_encoder = GoalPoseEncoder(
@@ -333,10 +345,12 @@ def forward(self, x, mask):
 
 
 class EgoEncoder(nn.Module):
-    def __init__(self, time_len, drop_path_rate, hidden_dim, depth):
+    def __init__(
+        self, time_len, drop_path_rate, hidden_dim, depth, mixer_hidden_dim, mixer_token_dim
+    ):
         super().__init__()
-        tokens_mlp_dim = 64
-        channels_mlp_dim = 128
+        tokens_mlp_dim = mixer_token_dim
+        channels_mlp_dim = mixer_hidden_dim
 
         self._hidden_dim = hidden_dim
 
@@ -396,10 +410,12 @@ def forward(self, x):
 
 
 class NeighborEncoder(nn.Module):
-    def __init__(self, time_len, drop_path_rate, hidden_dim, depth):
+    def __init__(
+        self, time_len, drop_path_rate, hidden_dim, depth, mixer_hidden_dim, mixer_token_dim
+    ):
         super().__init__()
-        tokens_mlp_dim = 64
-        channels_mlp_dim = 128
+        tokens_mlp_dim = mixer_token_dim
+        channels_mlp_dim = mixer_hidden_dim
 
         self._hidden_dim = hidden_dim
 
@@ -536,10 +552,19 @@ def forward(self, x):
 
 
 class LaneEncoder(nn.Module):
-    def __init__(self, lane_len, class_type, drop_path_rate, hidden_dim, depth):
+    def __init__(
+        self,
+        lane_len,
+        class_type,
+        drop_path_rate,
+        hidden_dim,
+        depth,
+        mixer_hidden_dim,
+        mixer_token_dim,
+    ):
         super().__init__()
-        tokens_mlp_dim = 64
-        channels_mlp_dim = 128
+        tokens_mlp_dim = mixer_token_dim
+        channels_mlp_dim = mixer_hidden_dim
 
         assert class_type in [CLASS_TYPE_LANE, CLASS_TYPE_ROUTE], (
             "Invalid class type for LaneEncoder"
@@ -640,11 +665,21 @@ def forward(self, x, speed_limit, has_speed_limit):
 
 
 class LineEncoder(nn.Module):
-    def __init__(self, line_len, class_type, drop_path_rate, hidden_dim, depth, point_dim=2):
+    def __init__(
+        self,
+        line_len,
+        class_type,
+        drop_path_rate,
+        hidden_dim,
+        depth,
+        mixer_hidden_dim,
+        mixer_token_dim,
+        point_dim,
+    ):
         super().__init__()
         self._class_type = class_type
-        tokens_mlp_dim = 64
-        channels_mlp_dim = 128
+        tokens_mlp_dim = mixer_token_dim
+        channels_mlp_dim = mixer_hidden_dim
 
         self._line_len = line_len
 

diffusion_planner/diffusion_planner/train_epoch.py

@@ -41,6 +41,8 @@ def train_epoch(data_loader, model, optimizer, args, ema, aug: StatePerturbation
 
     model.train()
 
+    device_type = "cuda" if "cuda" in str(args.device) else "cpu"
+
     if args.ddp:
         torch.cuda.synchronize()
 
@@ -69,15 +71,17 @@ def train_epoch(data_loader, model, optimizer, args, ema, aug: StatePerturbation
         # call the model
         optimizer.zero_grad()
 
-        loss = compute_training_loss(model, inputs, (ego_future, neighbors_future, mask), args)
-
-        loss["loss"] = (
-            args.alpha_neighbor_loss * loss["neighbor_prediction_loss"]
-            + args.alpha_planning_loss * loss["ego_planning_loss"]
-            + loss["turn_indicator_loss"]
-            + args.coeff_road_border_loss * loss["road_border_loss"]
-            + args.coeff_neighbor_collision_loss * loss["neighbor_collision_loss"]
-        )
+        # bf16 keeps the fp32 exponent range, so no GradScaler is needed.
+        with torch.autocast(device_type, dtype=torch.bfloat16, enabled=args.use_amp):
+            loss = compute_training_loss(model, inputs, (ego_future, neighbors_future, mask), args)
+
+            loss["loss"] = (
+                args.alpha_neighbor_loss * loss["neighbor_prediction_loss"]
+                + args.alpha_planning_loss * loss["ego_planning_loss"]
+                + loss["turn_indicator_loss"]
+                + args.coeff_road_border_loss * loss["road_border_loss"]
+                + args.coeff_neighbor_collision_loss * loss["neighbor_collision_loss"]
+            )
 
         # loss backward
         loss["loss"].backward()