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 <vibe@mistral.ai>
This commit is contained in:
2026-08-01 17:47:45 +08:00
parent a97b4da1ad
commit 479c2b1488
3 changed files with 25 additions and 16 deletions
+14 -10
View File
@@ -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}")