Compare commits
4 Commits
cc5dedc3fe
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| 479c2b1488 | |||
| a97b4da1ad | |||
| ce70d932d3 | |||
| 134bd2c3dc |
@@ -1,56 +1,49 @@
|
|||||||
# UZH-FPV Velocity Prediction
|
# UZH-FPV Velocity Prediction
|
||||||
|
|
||||||
从 DAVIS 事件相机灰度图像序列中预测**机体速度**(body-frame forward/lateral velocity)。
|
从 DAVIS 事件相机灰度图像序列预测机体速度(body-frame forward/lateral velocity)。
|
||||||
|
|
||||||
## 项目结构
|
## 项目结构
|
||||||
|
|
||||||
```
|
```
|
||||||
uzh_fpv/
|
uzh_fpv/
|
||||||
├── AGENTS.md # ← 本文件
|
├── AGENTS.md
|
||||||
├── requirements.txt # Python 依赖
|
├── requirements.txt
|
||||||
├── DATASET_FORMAT.md # 数据集格式详细说明
|
├── DATASET_FORMAT.md
|
||||||
├── rosbag2wds.py # ROS bag → WebDataset shard 转换脚本
|
├── rosbag2wds.py # ROS bag → WebDataset shard
|
||||||
├── batch_convert.sh # 批量转换脚本
|
├── batch_convert.sh
|
||||||
├── dataset/ # 数据集(.gitignore 忽略)
|
├── dataset/ # 数据集(.gitignore)
|
||||||
│ └── <scene_name>/
|
│ └── <scene_name>/
|
||||||
│ ├── shard_0000.tar # WebDataset shard(图像+GT)
|
│ ├── shard_0000.tar
|
||||||
│ ├── imu_sequence.npz # 完整 IMU 序列
|
│ ├── imu_sequence.npz
|
||||||
│ └── metadata.json # 元信息
|
│ └── metadata.json
|
||||||
├── src/
|
├── src/
|
||||||
│ ├── event_utils.py # EventProcessor: 帧间亮度变化 → 模拟事件帧
|
│ ├── event_utils.py # 帧间亮度变化 → 模拟事件帧
|
||||||
│ └── velocity_prediction/ # 主项目代码
|
│ └── velocity_prediction/ # 主项目代码
|
||||||
│ ├── __init__.py # 模块说明
|
│ ├── config.py
|
||||||
│ ├── config.py # 路径、模型架构、训练超参数
|
│ ├── utils.py
|
||||||
│ ├── utils.py # 四元数运算(torch + numpy 封装)
|
│ ├── transforms.py
|
||||||
│ ├── transforms.py # 数据预处理管线
|
│ ├── dataset.py
|
||||||
│ ├── dataset.py # WebDataset 加载 + 序列采样
|
│ ├── model.py
|
||||||
│ ├── model.py # CNN + PoseMLP + GRU + Head
|
│ ├── train.py
|
||||||
│ ├── train.py # 训练循环
|
│ └── evaluate.py
|
||||||
│ └── evaluate.py # 评估 + 绘图
|
|
||||||
├── visualize/
|
├── visualize/
|
||||||
│ ├── __init__.py
|
│ └── visualize_dataset.py
|
||||||
│ └── visualize_dataset.py # 数据集可视化:叠加位姿信息并生成视频
|
├── checkpoints/ # 模型权重(.gitignore)
|
||||||
├── benchmark/
|
├── logs/ # TensorBoard 日志(.gitignore)
|
||||||
│ ├── __init__.py
|
└── videos/ # 可视化输出(.gitignore)
|
||||||
│ ├── config.py # 评估配置
|
|
||||||
│ ├── evaluate.py # 完整评估管线
|
|
||||||
│ └── benchmark.py # 统一评估入口
|
|
||||||
├── checkpoints/ # 模型权重(.gitignore 忽略)
|
|
||||||
├── logs/ # TensorBoard 日志(.gitignore 忽略)
|
|
||||||
└── videos/ # 可视化输出视频
|
|
||||||
```
|
```
|
||||||
|
|
||||||
## 运行环境
|
## 运行环境
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv run python -m <module> # 使用 uv 虚拟环境运行
|
uv run python -m <module>
|
||||||
```
|
```
|
||||||
|
|
||||||
依赖见 `requirements.txt`,核心依赖:PyTorch、WebDataset、OpenCV、NumPy、Matplotlib。
|
依赖:PyTorch, WebDataset, OpenCV, NumPy, Matplotlib。
|
||||||
|
|
||||||
## 数据集
|
## 数据集
|
||||||
|
|
||||||
UZH-FPV 数据集,由 DAVIS 事件相机采集。每个场景目录包含:
|
UZH-FPV 数据集,DAVIS 事件相机采集。每个场景目录:
|
||||||
|
|
||||||
| 文件 | 格式 | 内容 |
|
| 文件 | 格式 | 内容 |
|
||||||
|------|------|------|
|
|------|------|------|
|
||||||
@@ -58,7 +51,7 @@ UZH-FPV 数据集,由 DAVIS 事件相机采集。每个场景目录包含:
|
|||||||
| `imu_sequence.npz` | NPZ | 完整 IMU 序列(加速度+角速度) |
|
| `imu_sequence.npz` | NPZ | 完整 IMU 序列(加速度+角速度) |
|
||||||
| `metadata.json` | JSON | 场景元信息 |
|
| `metadata.json` | JSON | 场景元信息 |
|
||||||
|
|
||||||
shard 中每个样本的字段:
|
shard 样本字段:
|
||||||
|
|
||||||
| Key | 类型 | 说明 |
|
| Key | 类型 | 说明 |
|
||||||
|-----|------|------|
|
|-----|------|------|
|
||||||
@@ -67,8 +60,6 @@ shard 中每个样本的字段:
|
|||||||
| `pose` | float32[7] | `[x, y, z, qx, qy, qz, qw]` 世界→机体四元数 |
|
| `pose` | float32[7] | `[x, y, z, qx, qy, qz, qw]` 世界→机体四元数 |
|
||||||
| `vel` | float32[6] | `[vx, vy, vz, wx, wy, wz]` 世界线速度 + 角速度 |
|
| `vel` | float32[6] | `[vx, vy, vz, wx, wy, wz]` 世界线速度 + 角速度 |
|
||||||
|
|
||||||
坐标系:z 轴与重力对齐(水平坐标系)。
|
|
||||||
|
|
||||||
### 场景列表
|
### 场景列表
|
||||||
|
|
||||||
| 场景 | 帧数 | 类型 |
|
| 场景 | 帧数 | 类型 |
|
||||||
@@ -78,62 +69,11 @@ shard 中每个样本的字段:
|
|||||||
| outdoor_forward_1/3/5 | 907~13299 | 室外前飞 |
|
| outdoor_forward_1/3/5 | 907~13299 | 室外前飞 |
|
||||||
| outdoor_45_1 | 799 | 室外 45° 飞行 |
|
| outdoor_45_1 | 799 | 室外 45° 飞行 |
|
||||||
|
|
||||||
## 模型
|
|
||||||
|
|
||||||
### 架构
|
|
||||||
|
|
||||||
```
|
|
||||||
Event frame (1, 240, 320) ──► CNN (4 Conv+Pool+GAP, 256-d)
|
|
||||||
│
|
|
||||||
Body up (3,) ──► PoseMLP (3→32→64, 64-d) ────────────────────────┤
|
|
||||||
│
|
|
||||||
concat (320-d) ← per-frame
|
|
||||||
│
|
|
||||||
GRU (hidden=128)
|
|
||||||
│
|
|
||||||
Head MLP (128→64→2)
|
|
||||||
│
|
|
||||||
[v_right, v_forward]
|
|
||||||
```
|
|
||||||
|
|
||||||
**注意**:当前 CNN 编码器被禁用(输出全零),模型仅依赖 `PoseMLP + GRU + Head`。
|
|
||||||
|
|
||||||
### 输入
|
|
||||||
|
|
||||||
- `events`: `(B, S, 1, H, W)` — 模拟事件帧,值域 `{-1, 0, +1}`
|
|
||||||
- `tilt`: `(B, S, 3)` — body up 向量(世界 up 旋转到机体坐标系),仅含 pitch/roll,不含 yaw,单位向量
|
|
||||||
|
|
||||||
### 输出
|
|
||||||
|
|
||||||
- `v_body`: `(B, 2)` — 机体坐标系 `[v_right, v_forward]` 速度 (m/s)
|
|
||||||
|
|
||||||
### 数据预处理管线
|
|
||||||
|
|
||||||
```
|
|
||||||
shard_*.tar → DecodeSample → SimulateEvents → ComputeTilt → ComputeBodyVelocity → NormalizeVelocity
|
|
||||||
```
|
|
||||||
|
|
||||||
1. **DecodeSample**: JPEG → 灰度图 uint8 (H,W);bytes → float32 数组
|
|
||||||
2. **SimulateEvents**: 帧间亮度变化 → 二值事件帧 `{-1, 0, +1}`
|
|
||||||
3. **ComputeTilt**: 四元数 (world→odom) → 应用 R_odom_to_body → 旋转 world-up [0,0,1] → body up 向量 (3,)
|
|
||||||
4. **ComputeBodyVelocity**: 世界速度 → 应用 R_odom_to_body → yaw 补偿(仅去除偏航,保留 tilt)→ 水平面 `[v_right, v_forward]`
|
|
||||||
5. **NormalizeVelocity**: 归一化
|
|
||||||
|
|
||||||
### 训练配置
|
|
||||||
|
|
||||||
- seq_len=8, batch_size=32, epochs=100
|
|
||||||
- lr=1e-3, AdamW, StepLR (step=30, gamma=0.5)
|
|
||||||
- Loss: MSELoss
|
|
||||||
- 训练/验证/测试场景见 `config.py`
|
|
||||||
|
|
||||||
## 关键命令
|
## 关键命令
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# 训练
|
# 训练(GPU 优先 cuda:7)
|
||||||
uv run python -m src.velocity_prediction.train --device cuda:0
|
uv run python -m src.velocity_prediction.train --device cuda:7
|
||||||
|
|
||||||
# 评估
|
|
||||||
uv run python -m src.velocity_prediction.evaluate --checkpoint checkpoints/best.pt
|
|
||||||
|
|
||||||
# 数据集可视化(单场景)
|
# 数据集可视化(单场景)
|
||||||
uv run python -m visualize.visualize_dataset --scene indoor_forward_3 --output videos/scene.mp4
|
uv run python -m visualize.visualize_dataset --scene indoor_forward_3 --output videos/scene.mp4
|
||||||
@@ -143,49 +83,4 @@ uv run python -m visualize.visualize_dataset --all --output videos/
|
|||||||
|
|
||||||
# 数据集可视化(实时显示)
|
# 数据集可视化(实时显示)
|
||||||
uv run python -m visualize.visualize_dataset --scene indoor_forward_3 --show
|
uv run python -m visualize.visualize_dataset --scene indoor_forward_3 --show
|
||||||
|
|
||||||
# Benchmark 评估
|
|
||||||
uv run python -m benchmark.benchmark --checkpoint checkpoints/best.pt
|
|
||||||
```
|
```
|
||||||
|
|
||||||
## 可视化说明
|
|
||||||
|
|
||||||
`visualize/visualize_dataset.py` 在每帧图像上叠加:
|
|
||||||
|
|
||||||
- 帧号、时间戳、世界坐标位置
|
|
||||||
- 欧拉角 `[roll, pitch, yaw]`(从 body 四元数计算)
|
|
||||||
- Body up 向量 `[x, y, z]`
|
|
||||||
- 机体速度 `v_body [forward, lateral]`
|
|
||||||
- 世界速度 `v_world [vx, vy, vz]`
|
|
||||||
- 机体坐标系三轴箭头(左下角)
|
|
||||||
- 机体速度方向箭头(图像中心)
|
|
||||||
|
|
||||||
## 关键约定
|
|
||||||
|
|
||||||
- 四元数格式:`[x, y, z, w]`(不是 `[w, x, y, z]`)
|
|
||||||
- GT 四元数表示 **world→odom**(不是 world→body),通过静态 R_odom_to_body 校正
|
|
||||||
- Body 坐标系(ROS 右手系):`body_x=右, body_y=前, body_z=上`
|
|
||||||
- R_odom_to_body = R_y(45°) @ R_x(90°):先绕 odom_x 转 +90°,再绕 odom_y 转 +45°
|
|
||||||
- 速度归一化统计量:待重新计算
|
|
||||||
- 模型预测 `[v_right, v_forward]`(右向和前向速度)
|
|
||||||
- 所有代码在项目根目录下以 `uv run python -m <module>` 运行
|
|
||||||
- GPU 优先使用 `cuda:7`,训练时添加 `--device cuda:7`
|
|
||||||
|
|
||||||
## 已知问题
|
|
||||||
|
|
||||||
### 1. 滑窗跨 shard 边界
|
|
||||||
|
|
||||||
`dataset.py` 中滑窗实现基于 WebDataset 串联后的连续流,不感知 shard 边界。当样本恰好处于 shard 末尾时,序列会跨越到下一个 shard 的起始帧。
|
|
||||||
|
|
||||||
- 影响:每个 shard 边界处约有 `seq_len` 个序列包含跨 shard 样本(占总数 <1%)
|
|
||||||
- 修复思路:在 `_sliding_window_fn` 中注入 shard 边界标记,遇到边界时清空缓冲区
|
|
||||||
- 严重程度:低。若 shard 内帧数远大于 seq_len,可忽略
|
|
||||||
|
|
||||||
### 2. `SimulateEvents` 跨 shard 状态残留
|
|
||||||
|
|
||||||
`EventProcessor` 内部维护 `_prev_frame` 用于帧差计算。跨 shard 时,新 shard 的第一帧会与上一个 shard 最后一帧计算差,产生错误的事件帧。
|
|
||||||
|
|
||||||
- 影响:每个 shard 的第 1 帧事件帧错误,涉及该帧的所有滑窗序列均受影响
|
|
||||||
- 每 shard 错误帧数:1 帧(加上滑窗放大,约 `seq_len` 个序列各包含此帧)
|
|
||||||
- 修复:在 shard 边界处调用 `EventProcessor.reset()`。需在 `_build_pipeline` 中插入边界信号或改用按 shard 独立处理的方案
|
|
||||||
- 严重程度:低。每 shard 仅 1 帧,训练数据量大时可忽略
|
|
||||||
|
|||||||
@@ -39,10 +39,10 @@ VAL_SCENES = [
|
|||||||
# "indoor_forward_3", "indoor_forward_9", "indoor_forward_10", # Easy
|
# "indoor_forward_3", "indoor_forward_9", "indoor_forward_10", # Easy
|
||||||
]
|
]
|
||||||
TEST_SCENES = [
|
TEST_SCENES = [
|
||||||
"indoor_forward_9",
|
# "indoor_forward_9",
|
||||||
# "indoor_forward_9","indoor_forward_3",
|
# "indoor_forward_9","indoor_forward_3",
|
||||||
# "indoor_forward_7", # Hard 室内
|
# "indoor_forward_7", # Hard 室内
|
||||||
# "outdoor_forward_1", # Easy 室外
|
"outdoor_forward_1", # Easy 室外
|
||||||
# "outdoor_forward_5" # Hard 室外
|
# "outdoor_forward_5" # Hard 室外
|
||||||
# "indoor_forward_3", "indoor_forward_9", "indoor_forward_10", # Easy
|
# "indoor_forward_3", "indoor_forward_9", "indoor_forward_10", # Easy
|
||||||
]
|
]
|
||||||
@@ -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
|
||||||
|
|||||||
@@ -166,7 +166,11 @@ def main():
|
|||||||
ckpt = torch.load(args.checkpoint, map_location="cpu")
|
ckpt = torch.load(args.checkpoint, map_location="cpu")
|
||||||
model.load_state_dict(ckpt["model_state_dict"])
|
model.load_state_dict(ckpt["model_state_dict"])
|
||||||
model.to(device)
|
model.to(device)
|
||||||
|
|
||||||
|
# Restore event threshold from checkpoint, fall back to config default
|
||||||
|
event_threshold = ckpt.get("event_threshold", train_cfg.event_threshold)
|
||||||
print(f"Loaded checkpoint from {args.checkpoint} (epoch={ckpt.get('epoch', '?')})")
|
print(f"Loaded checkpoint from {args.checkpoint} (epoch={ckpt.get('epoch', '?')})")
|
||||||
|
print(f"Event threshold: {event_threshold} (checkpoint={ckpt.get('event_threshold', 'not saved')}, config={train_cfg.event_threshold})")
|
||||||
|
|
||||||
# Evaluate each scene independently → NaN gaps prevent plot mixing
|
# Evaluate each scene independently → NaN gaps prevent plot mixing
|
||||||
from src.velocity_prediction.config import TEST_SCENES
|
from src.velocity_prediction.config import TEST_SCENES
|
||||||
@@ -180,7 +184,7 @@ def main():
|
|||||||
stride=1,
|
stride=1,
|
||||||
batch_size=1,
|
batch_size=1,
|
||||||
num_workers=0, # strict temporal order
|
num_workers=0, # strict temporal order
|
||||||
event_threshold=train_cfg.event_threshold,
|
event_threshold=event_threshold,
|
||||||
event_use_log=train_cfg.event_use_log,
|
event_use_log=train_cfg.event_use_log,
|
||||||
)
|
)
|
||||||
results = evaluate_stateful(model, loader, device)
|
results = evaluate_stateful(model, loader, device)
|
||||||
|
|||||||
@@ -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,14 +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]
|
||||||
# cnn_feat = events.new_zeros(B, S, self.cnn.out_dim) # 全零替代
|
|
||||||
|
|
||||||
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)
|
||||||
@@ -174,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)
|
||||||
@@ -195,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}")
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ Usage:
|
|||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import os
|
import os
|
||||||
|
import random
|
||||||
import time
|
import time
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -39,9 +40,9 @@ def train_one_epoch(
|
|||||||
log_interval: int = 50,
|
log_interval: int = 50,
|
||||||
global_step: int = 0,
|
global_step: int = 0,
|
||||||
use_amp: bool = True,
|
use_amp: bool = True,
|
||||||
h_state: torch.Tensor = None,
|
) -> tuple[float, int]:
|
||||||
) -> tuple[float, int, torch.Tensor]:
|
"""Train for one epoch (stateless — each batch starts from zero hidden state).
|
||||||
"""Train for one epoch. Returns (avg_loss, updated_global_step, h_state_final)."""
|
Returns (avg_loss, updated_global_step)."""
|
||||||
model.train()
|
model.train()
|
||||||
total_loss = 0.0
|
total_loss = 0.0
|
||||||
num_batches = 0
|
num_batches = 0
|
||||||
@@ -52,15 +53,11 @@ def train_one_epoch(
|
|||||||
tilt = batch["tilt"].to(device) # (B, S, 3)
|
tilt = batch["tilt"].to(device) # (B, S, 3)
|
||||||
target = batch["v_body_target"].to(device) # (B, S, 2)
|
target = batch["v_body_target"].to(device) # (B, S, 2)
|
||||||
|
|
||||||
# Drop carried state if batch size changed (partial last batch / new scene)
|
# Stateless: each batch starts from zero hidden state
|
||||||
if h_state is not None and h_state.shape[1] != events.shape[0]:
|
|
||||||
h_state = None
|
|
||||||
|
|
||||||
# Per-step supervision over the whole sequence — TBPTT hidden state
|
|
||||||
with torch.amp.autocast(device.type, enabled=use_amp):
|
with torch.amp.autocast(device.type, enabled=use_amp):
|
||||||
pred_seq, h_new = model(events, tilt, h_state) # (B, S, 2), (L, B, H)
|
pred_seq, _ = model(events, tilt, None) # (B, S, 2), h=None
|
||||||
loss_per_step = criterion(pred_seq, target) # (B, S, 2)
|
loss_per_step = criterion(pred_seq, target) # (B, S, 2)
|
||||||
loss_per_step = loss_per_step.mean(-1) # (B, S)
|
loss_per_step = loss_per_step.mean(-1) # (B, S)
|
||||||
loss = loss_per_step.mean()
|
loss = loss_per_step.mean()
|
||||||
|
|
||||||
optimizer.zero_grad()
|
optimizer.zero_grad()
|
||||||
@@ -68,9 +65,6 @@ def train_one_epoch(
|
|||||||
scaler.step(optimizer)
|
scaler.step(optimizer)
|
||||||
scaler.update()
|
scaler.update()
|
||||||
|
|
||||||
# Detach hidden state for next batch — gradient only back to current seq_len
|
|
||||||
h_state = h_new.detach()
|
|
||||||
|
|
||||||
total_loss += loss.item()
|
total_loss += loss.item()
|
||||||
num_batches += 1
|
num_batches += 1
|
||||||
global_step += 1
|
global_step += 1
|
||||||
@@ -87,7 +81,7 @@ def train_one_epoch(
|
|||||||
|
|
||||||
avg_loss = total_loss / max(num_batches, 1)
|
avg_loss = total_loss / max(num_batches, 1)
|
||||||
print(f" Epoch {epoch} | Avg Loss: {avg_loss:.6f}")
|
print(f" Epoch {epoch} | Avg Loss: {avg_loss:.6f}")
|
||||||
return avg_loss, global_step, h_state
|
return avg_loss, global_step
|
||||||
|
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
@@ -150,15 +144,12 @@ def main():
|
|||||||
print(f"Model parameters: {total_params:,} ({total_params/1e6:.3f} M)")
|
print(f"Model parameters: {total_params:,} ({total_params/1e6:.3f} M)")
|
||||||
print(f"AMP: {'enabled' if use_amp else 'disabled'}")
|
print(f"AMP: {'enabled' if use_amp else 'disabled'}")
|
||||||
|
|
||||||
# Data loaders — TBPTT training requires strict temporal order
|
# ── Randomised sequence lengths for curriculum ──────────────
|
||||||
train_loader = create_tbptt_loader(
|
SEQ_LEN_CHOICES = [64, 96, 128, 160, 192]
|
||||||
seq_len=train_cfg.seq_len,
|
|
||||||
batch_size=train_cfg.batch_size,
|
# Validation uses a fixed seq_len for consistent comparison
|
||||||
event_threshold=event_threshold,
|
|
||||||
event_use_log=train_cfg.event_use_log,
|
|
||||||
)
|
|
||||||
val_loader = create_val_loader(
|
val_loader = create_val_loader(
|
||||||
seq_len=train_cfg.seq_len,
|
seq_len=128,
|
||||||
stride=train_cfg.sliding_window_stride,
|
stride=train_cfg.sliding_window_stride,
|
||||||
batch_size=train_cfg.batch_size,
|
batch_size=train_cfg.batch_size,
|
||||||
num_workers=train_cfg.num_workers,
|
num_workers=train_cfg.num_workers,
|
||||||
@@ -200,6 +191,10 @@ def main():
|
|||||||
global_step = ckpt.get("global_step", 0)
|
global_step = ckpt.get("global_step", 0)
|
||||||
best_val_loss = ckpt.get("best_val_loss", float("inf"))
|
best_val_loss = ckpt.get("best_val_loss", float("inf"))
|
||||||
run_id = ckpt.get("run_id", None)
|
run_id = ckpt.get("run_id", None)
|
||||||
|
# Restore event threshold from checkpoint if present
|
||||||
|
if "event_threshold" in ckpt:
|
||||||
|
event_threshold = ckpt["event_threshold"]
|
||||||
|
print(f"Event threshold restored from checkpoint: {event_threshold}")
|
||||||
|
|
||||||
print(f"Resumed from checkpoint: {ckpt_path}")
|
print(f"Resumed from checkpoint: {ckpt_path}")
|
||||||
print(f" Resumed epoch={ckpt.get('epoch', '?')}, global_step={global_step}, "
|
print(f" Resumed epoch={ckpt.get('epoch', '?')}, global_step={global_step}, "
|
||||||
@@ -216,20 +211,28 @@ def main():
|
|||||||
ckpt_dir.mkdir(parents=True, exist_ok=True)
|
ckpt_dir.mkdir(parents=True, exist_ok=True)
|
||||||
writer = SummaryWriter(log_dir=str(log_dir))
|
writer = SummaryWriter(log_dir=str(log_dir))
|
||||||
|
|
||||||
print(f" seq_len={train_cfg.seq_len}, batch_size={train_cfg.batch_size}")
|
print(f" seq_len=random({min(SEQ_LEN_CHOICES)}~{max(SEQ_LEN_CHOICES)}), "
|
||||||
|
f"batch_size={train_cfg.batch_size}")
|
||||||
print(f" lr={train_cfg.lr}, weight_decay={train_cfg.weight_decay}")
|
print(f" lr={train_cfg.lr}, weight_decay={train_cfg.weight_decay}")
|
||||||
print(f" log_dir={log_dir}, checkpoint_dir={ckpt_dir}\n")
|
print(f" log_dir={log_dir}, checkpoint_dir={ckpt_dir}\n")
|
||||||
|
|
||||||
for epoch in range(start_epoch, train_cfg.epochs + 1):
|
for epoch in range(start_epoch, train_cfg.epochs + 1):
|
||||||
epoch_start = time.time()
|
epoch_start = time.time()
|
||||||
|
|
||||||
# Reset GRU hidden state at epoch start — each epoch begins with h=0
|
# Random seq_len per epoch — GRU learns to handle varying temporal horizons
|
||||||
train_loss, global_step, _ = train_one_epoch(
|
seq_len = random.choice(SEQ_LEN_CHOICES)
|
||||||
|
train_loader = create_tbptt_loader(
|
||||||
|
seq_len=seq_len,
|
||||||
|
batch_size=train_cfg.batch_size,
|
||||||
|
event_threshold=event_threshold,
|
||||||
|
event_use_log=train_cfg.event_use_log,
|
||||||
|
)
|
||||||
|
|
||||||
|
train_loss, global_step = train_one_epoch(
|
||||||
model, train_loader, optimizer, criterion, scaler, device, epoch, writer,
|
model, train_loader, optimizer, criterion, scaler, device, epoch, writer,
|
||||||
log_interval=train_cfg.log_interval,
|
log_interval=train_cfg.log_interval,
|
||||||
global_step=global_step,
|
global_step=global_step,
|
||||||
use_amp=use_amp,
|
use_amp=use_amp,
|
||||||
h_state=None,
|
|
||||||
)
|
)
|
||||||
val_loss = validate(model, val_loader, criterion, device, use_amp=use_amp)
|
val_loss = validate(model, val_loader, criterion, device, use_amp=use_amp)
|
||||||
scheduler.step()
|
scheduler.step()
|
||||||
@@ -238,6 +241,7 @@ def main():
|
|||||||
current_lr = scheduler.get_last_lr()[0]
|
current_lr = scheduler.get_last_lr()[0]
|
||||||
|
|
||||||
print(f"Epoch {epoch:3d}/{train_cfg.epochs} | "
|
print(f"Epoch {epoch:3d}/{train_cfg.epochs} | "
|
||||||
|
f"seq_len={seq_len:3d} | "
|
||||||
f"Train Loss: {train_loss:.6f} | Val Loss: {val_loss:.6f} | "
|
f"Train Loss: {train_loss:.6f} | Val Loss: {val_loss:.6f} | "
|
||||||
f"LR: {current_lr:.2e} | Time: {epoch_time:.1f}s")
|
f"LR: {current_lr:.2e} | Time: {epoch_time:.1f}s")
|
||||||
|
|
||||||
@@ -253,6 +257,7 @@ def main():
|
|||||||
"model_state_dict": model.state_dict(),
|
"model_state_dict": model.state_dict(),
|
||||||
"optimizer_state_dict": optimizer.state_dict(),
|
"optimizer_state_dict": optimizer.state_dict(),
|
||||||
"scheduler_state_dict": scheduler.state_dict(),
|
"scheduler_state_dict": scheduler.state_dict(),
|
||||||
|
"event_threshold": event_threshold,
|
||||||
"global_step": global_step,
|
"global_step": global_step,
|
||||||
"best_val_loss": best_val_loss,
|
"best_val_loss": best_val_loss,
|
||||||
"run_id": run_id,
|
"run_id": run_id,
|
||||||
@@ -269,6 +274,7 @@ def main():
|
|||||||
"model_state_dict": model.state_dict(),
|
"model_state_dict": model.state_dict(),
|
||||||
"optimizer_state_dict": optimizer.state_dict(),
|
"optimizer_state_dict": optimizer.state_dict(),
|
||||||
"scheduler_state_dict": scheduler.state_dict(),
|
"scheduler_state_dict": scheduler.state_dict(),
|
||||||
|
"event_threshold": event_threshold,
|
||||||
"global_step": global_step,
|
"global_step": global_step,
|
||||||
"best_val_loss": best_val_loss,
|
"best_val_loss": best_val_loss,
|
||||||
"run_id": run_id,
|
"run_id": run_id,
|
||||||
|
|||||||
@@ -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