diff --git a/diffusion_planner/diffusion_planner/config/model_config.py b/diffusion_planner/diffusion_planner/config/model_config.py index ca7e9b443..59236ac4a 100644 --- a/diffusion_planner/diffusion_planner/config/model_config.py +++ b/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 diff --git a/diffusion_planner/diffusion_planner/config/train_config.py b/diffusion_planner/diffusion_planner/config/train_config.py index a7280d1a1..914dbbf1a 100644 --- a/diffusion_planner/diffusion_planner/config/train_config.py +++ b/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 ) diff --git a/diffusion_planner/diffusion_planner/model/module/encoder.py b/diffusion_planner/diffusion_planner/model/module/encoder.py index ac63ce789..9a3ae9558 100644 --- a/diffusion_planner/diffusion_planner/model/module/encoder.py +++ b/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,6 +92,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, ) self.route_encoder = LaneEncoder( config.route_len, @@ -95,6 +101,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, ) self.polygon_encoder = LineEncoder( config.polygon_len, @@ -102,6 +110,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 + 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 diff --git a/diffusion_planner/diffusion_planner/train_epoch.py b/diffusion_planner/diffusion_planner/train_epoch.py index d71fe22c9..a762dab84 100644 --- a/diffusion_planner/diffusion_planner/train_epoch.py +++ b/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()