From 479c2b1488ca369b960c4d567204f9c6acdbea05 Mon Sep 17 00:00:00 2001 From: CaoWangrenbo Date: Sat, 1 Aug 2026 17:47:45 +0800 Subject: [PATCH] refactor: strided conv encoder with 80x60 input resolution - CNNEncoder: stride=2 convs replace Conv2d+MaxPool2d pattern - 3 layers (32,64,128) instead of 4 (32,64,128,256), GRU input 192 - DecodeSample resizes grayscale frames to 80x60 via INTER_AREA - Model params: 227K (was 1.5M), input 80x60 (was 320x240) Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- src/velocity_prediction/config.py | 8 ++++---- src/velocity_prediction/model.py | 24 ++++++++++++++---------- src/velocity_prediction/transforms.py | 9 +++++++-- 3 files changed, 25 insertions(+), 16 deletions(-) diff --git a/src/velocity_prediction/config.py b/src/velocity_prediction/config.py index f8469f8..ea52a5f 100644 --- a/src/velocity_prediction/config.py +++ b/src/velocity_prediction/config.py @@ -53,9 +53,9 @@ TEST_SCENES = [ @dataclass class CNNConfig: in_channels: int = 1 - channels: tuple = (32, 64, 128, 256) # per-layer output channels + channels: tuple = (32, 64, 128) # per-layer output channels kernel_size: int = 3 - pool_size: int = 2 + stride: int = 2 # strided conv replaces conv+pool use_bn: bool = True @@ -68,7 +68,7 @@ class PoseMLPConfig: @dataclass class GRUConfig: - input_size: int = 320 # CNN(256) + PoseMLP(64) + input_size: int = 192 # CNN(128) + PoseMLP(64) hidden_size: int = 128 num_layers: int = 1 dropout: float = 0.0 @@ -104,7 +104,7 @@ class TrainConfig: seed: int = 42 # Sliding window: stride=1 → full overlap, stride=seq_len → non-overlapping - sliding_window_stride: int = 1 + sliding_window_stride: int = 64 # Event simulation event_threshold: float = 0.1 diff --git a/src/velocity_prediction/model.py b/src/velocity_prediction/model.py index 9d5d169..6f39835 100644 --- a/src/velocity_prediction/model.py +++ b/src/velocity_prediction/model.py @@ -14,7 +14,10 @@ from src.velocity_prediction.config import model_cfg class CNNEncoder(nn.Module): """ - 4-layer ConvNet with BatchNorm, ReLU, MaxPool, ending with Global Avg Pool. + Strided-convolution encoder with BatchNorm and LeakyReLU. + + Each layer uses stride=2 to downsample, replacing the traditional + conv+pool pattern. Ends with Global Avg Pool. Input: (B, S, 1, H, W) — processed per-frame (flattened to (B*S, 1, H, W)) Output: (B, S, C_out) — per-frame feature vectors @@ -24,15 +27,15 @@ class CNNEncoder(nn.Module): super().__init__() channels = cfg.channels in_ch = cfg.in_channels + stride = cfg.stride layers = [] for out_ch in channels: layers.extend([ - nn.Conv2d(in_ch, out_ch, kernel_size=cfg.kernel_size, padding=cfg.kernel_size // 2), + nn.Conv2d(in_ch, out_ch, kernel_size=cfg.kernel_size, + stride=stride, padding=cfg.kernel_size // 2), nn.BatchNorm2d(out_ch) if cfg.use_bn else nn.Identity(), - nn.Identity(), nn.LeakyReLU(inplace=True), - nn.MaxPool2d(cfg.pool_size), ]) in_ch = out_ch @@ -102,7 +105,7 @@ class VelocityPredictionModel(nn.Module): self.cnn = CNNEncoder(cnn_cfg) self.pose_mlp = PoseMLP(pose_cfg) - fused_dim = self.cnn.out_dim + self.pose_mlp.out_dim # 256 + 64 = 320 + fused_dim = self.cnn.out_dim + self.pose_mlp.out_dim # 128 + 64 = 192 self.gru = nn.GRU( input_size=fused_dim, @@ -139,13 +142,13 @@ class VelocityPredictionModel(nn.Module): B, S = events.shape[:2] # Per-frame encoding - cnn_feat = self.cnn(events) # (B, S, 256) + cnn_feat = self.cnn(events) # (B, S, 128) # B, S = events.shape[:2] pose_feat = self.pose_mlp(tilt) # (B, S, 64) # Fuse per frame - fused = torch.cat([cnn_feat, pose_feat], dim=-1) # (B, S, 320) + fused = torch.cat([cnn_feat, pose_feat], dim=-1) # (B, S, 192) # GRU temporal modelling — accepts external hidden state for TBPTT gru_out, h_new = self.gru(fused, h) # (B, S, 128), (num_layers, B, 128) @@ -173,9 +176,9 @@ class VelocityPredictionModel(nn.Module): v_body: (B, 2) body-frame [v_right, v_forward] h_new: (num_layers, B, hidden_size) """ - cnn_feat = self.cnn(events) # (B, 1, 256) + cnn_feat = self.cnn(events) # (B, 1, 128) pose_feat = self.pose_mlp(tilt) # (B, 1, 64) - fused = torch.cat([cnn_feat, pose_feat], dim=-1) # (B, 1, 320) + fused = torch.cat([cnn_feat, pose_feat], dim=-1) # (B, 1, 192) _, h_new = self.gru(fused, h) # (num_layers, B, 128) last_hidden = h_new[-1] # (B, 128) v_body = self.head(last_hidden) # (B, 2) @@ -194,10 +197,11 @@ if __name__ == "__main__": print(f"Total trainable parameters: {total:,} ({total/1e6:.3f} M)") # Forward pass test - B, S, H, W = 4, 8, 240, 320 + B, S, H, W = 4, 8, 60, 80 events = torch.randn(B, S, 1, H, W) tilt = torch.randn(B, S, 3) out, h = model(events, tilt) print(f"Input events: {events.shape}") print(f"Input tilt: {tilt.shape}") print(f"Output: {out.shape} (should be [4, 8, 2])") + print(f"CNN out_dim: {model.cnn.out_dim}") diff --git a/src/velocity_prediction/transforms.py b/src/velocity_prediction/transforms.py index 5f3aa39..5937786 100644 --- a/src/velocity_prediction/transforms.py +++ b/src/velocity_prediction/transforms.py @@ -20,11 +20,16 @@ from src.velocity_prediction.config import VELOCITY_MEAN, VELOCITY_STD class DecodeSample: - """Decode raw bytes from WebDataset tar entry into numpy arrays.""" + """Decode raw bytes from WebDataset tar entry into numpy arrays, resize to 80×60.""" + + def __init__(self, height: int = 60, width: int = 80): + self.height = height + self.width = width def __call__(self, sample: dict) -> dict: - # Image: JPEG bytes → grayscale uint8 (H, W) + # Image: JPEG bytes → grayscale uint8 (H, W) → resize img = cv2.imdecode(np.frombuffer(sample["jpg"], np.uint8), cv2.IMREAD_GRAYSCALE) + img = cv2.resize(img, (self.width, self.height), interpolation=cv2.INTER_AREA) # Timestamp ts = np.frombuffer(sample["ts"], dtype=np.float64).item()