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:
@@ -53,9 +53,9 @@ TEST_SCENES = [
|
|||||||
@dataclass
|
@dataclass
|
||||||
class CNNConfig:
|
class CNNConfig:
|
||||||
in_channels: int = 1
|
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
|
kernel_size: int = 3
|
||||||
pool_size: int = 2
|
stride: int = 2 # strided conv replaces conv+pool
|
||||||
use_bn: bool = True
|
use_bn: bool = True
|
||||||
|
|
||||||
|
|
||||||
@@ -68,7 +68,7 @@ class PoseMLPConfig:
|
|||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class GRUConfig:
|
class GRUConfig:
|
||||||
input_size: int = 320 # CNN(256) + PoseMLP(64)
|
input_size: int = 192 # CNN(128) + PoseMLP(64)
|
||||||
hidden_size: int = 128
|
hidden_size: int = 128
|
||||||
num_layers: int = 1
|
num_layers: int = 1
|
||||||
dropout: float = 0.0
|
dropout: float = 0.0
|
||||||
@@ -104,7 +104,7 @@ class TrainConfig:
|
|||||||
seed: int = 42
|
seed: int = 42
|
||||||
|
|
||||||
# Sliding window: stride=1 → full overlap, stride=seq_len → non-overlapping
|
# Sliding window: stride=1 → full overlap, stride=seq_len → non-overlapping
|
||||||
sliding_window_stride: int = 1
|
sliding_window_stride: int = 64
|
||||||
|
|
||||||
# Event simulation
|
# Event simulation
|
||||||
event_threshold: float = 0.1
|
event_threshold: float = 0.1
|
||||||
|
|||||||
@@ -14,7 +14,10 @@ from src.velocity_prediction.config import model_cfg
|
|||||||
|
|
||||||
class CNNEncoder(nn.Module):
|
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))
|
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
|
Output: (B, S, C_out) — per-frame feature vectors
|
||||||
@@ -24,15 +27,15 @@ class CNNEncoder(nn.Module):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
channels = cfg.channels
|
channels = cfg.channels
|
||||||
in_ch = cfg.in_channels
|
in_ch = cfg.in_channels
|
||||||
|
stride = cfg.stride
|
||||||
|
|
||||||
layers = []
|
layers = []
|
||||||
for out_ch in channels:
|
for out_ch in channels:
|
||||||
layers.extend([
|
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.BatchNorm2d(out_ch) if cfg.use_bn else nn.Identity(),
|
||||||
nn.Identity(),
|
|
||||||
nn.LeakyReLU(inplace=True),
|
nn.LeakyReLU(inplace=True),
|
||||||
nn.MaxPool2d(cfg.pool_size),
|
|
||||||
])
|
])
|
||||||
in_ch = out_ch
|
in_ch = out_ch
|
||||||
|
|
||||||
@@ -102,7 +105,7 @@ class VelocityPredictionModel(nn.Module):
|
|||||||
self.cnn = CNNEncoder(cnn_cfg)
|
self.cnn = CNNEncoder(cnn_cfg)
|
||||||
self.pose_mlp = PoseMLP(pose_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(
|
self.gru = nn.GRU(
|
||||||
input_size=fused_dim,
|
input_size=fused_dim,
|
||||||
@@ -139,13 +142,13 @@ class VelocityPredictionModel(nn.Module):
|
|||||||
B, S = events.shape[:2]
|
B, S = events.shape[:2]
|
||||||
|
|
||||||
# Per-frame encoding
|
# Per-frame encoding
|
||||||
cnn_feat = self.cnn(events) # (B, S, 256)
|
cnn_feat = self.cnn(events) # (B, S, 128)
|
||||||
# B, S = events.shape[:2]
|
# B, S = events.shape[:2]
|
||||||
|
|
||||||
pose_feat = self.pose_mlp(tilt) # (B, S, 64)
|
pose_feat = self.pose_mlp(tilt) # (B, S, 64)
|
||||||
|
|
||||||
# Fuse per frame
|
# 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 temporal modelling — accepts external hidden state for TBPTT
|
||||||
gru_out, h_new = self.gru(fused, h) # (B, S, 128), (num_layers, B, 128)
|
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]
|
v_body: (B, 2) body-frame [v_right, v_forward]
|
||||||
h_new: (num_layers, B, hidden_size)
|
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)
|
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)
|
_, h_new = self.gru(fused, h) # (num_layers, B, 128)
|
||||||
last_hidden = h_new[-1] # (B, 128)
|
last_hidden = h_new[-1] # (B, 128)
|
||||||
v_body = self.head(last_hidden) # (B, 2)
|
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)")
|
print(f"Total trainable parameters: {total:,} ({total/1e6:.3f} M)")
|
||||||
|
|
||||||
# Forward pass test
|
# 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)
|
events = torch.randn(B, S, 1, H, W)
|
||||||
tilt = torch.randn(B, S, 3)
|
tilt = torch.randn(B, S, 3)
|
||||||
out, h = model(events, tilt)
|
out, h = model(events, tilt)
|
||||||
print(f"Input events: {events.shape}")
|
print(f"Input events: {events.shape}")
|
||||||
print(f"Input tilt: {tilt.shape}")
|
print(f"Input tilt: {tilt.shape}")
|
||||||
print(f"Output: {out.shape} (should be [4, 8, 2])")
|
print(f"Output: {out.shape} (should be [4, 8, 2])")
|
||||||
|
print(f"CNN out_dim: {model.cnn.out_dim}")
|
||||||
|
|||||||
@@ -20,11 +20,16 @@ from src.velocity_prediction.config import VELOCITY_MEAN, VELOCITY_STD
|
|||||||
|
|
||||||
|
|
||||||
class DecodeSample:
|
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:
|
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.imdecode(np.frombuffer(sample["jpg"], np.uint8), cv2.IMREAD_GRAYSCALE)
|
||||||
|
img = cv2.resize(img, (self.width, self.height), interpolation=cv2.INTER_AREA)
|
||||||
|
|
||||||
# Timestamp
|
# Timestamp
|
||||||
ts = np.frombuffer(sample["ts"], dtype=np.float64).item()
|
ts = np.frombuffer(sample["ts"], dtype=np.float64).item()
|
||||||
|
|||||||
Reference in New Issue
Block a user