BEVFormer 训练实战

conda activate ad

# 只需要这些 —— 没有 mmcv、没有版本地狱
pip install torch==2.1.2 torchvision==0.16.2 \
    --index-url https://download.pytorch.org/whl/cu121
pip install tensorboard tqdm

# 验证
python -c "
import torch, torchvision
print('torch', torch.__version__)
print('cuda 可用:', torch.cuda.is_available())
print('设备:', torch.cuda.get_device_name(0))
print('算力: sm_%d%d' % torch.cuda.get_device_capability(0))
"
# 期望输出算力 sm_89，且 cuda 可用为 True

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MiniBEVFormer —— 从零实现的精简版 BEVFormer

与论文版的差异（诚实说明）:
    · 空间交叉注意力用「固定采样点」而非可变形偏移（机制相同）
    · 检测头用 CenterPoint 风格热力图，而非 DETR query（更易收敛）
    · 时序融合作为可选项
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision


# ==================== 图像骨干 ====================
class ImageBackbone(nn.Module):
    """ResNet 提取多相机图像特征"""

    def __init__(self, out_channels=256, pretrained=True):
        super().__init__()
        net = torchvision.models.resnet50(
            weights=torchvision.models.ResNet50_Weights.DEFAULT
            if pretrained else None)
        # 取到 layer3，输出 stride=16、通道 1024
        self.stem = nn.Sequential(
            net.conv1, net.bn1, net.relu, net.maxpool,
            net.layer1, net.layer2, net.layer3)
        self.proj = nn.Conv2d(1024, out_channels, 1)

    def forward(self, imgs):
        """imgs [B,N,3,H,W] → [B,N,C,H/16,W/16]"""
        B, N = imgs.shape[:2]
        x = imgs.flatten(0, 1)              # [B*N,3,H,W]
        x = self.stem(x)
        x = self.proj(x)
        _, C, h, w = x.shape
        return x.view(B, N, C, h, w)


# ==================== 空间交叉注意力（核心）====================
class SpatialCrossAttention(nn.Module):
    """
    BEV 格子 → 多相机图像 的特征采样

    这是 BEVFormer 论文第 4 节的核心机制:
      1. 每个 BEV 格子对应一根「柱子」，在柱子上取 Z 个高度点
      2. 把这些 3D 点用 lidar2img 投影到 6 路相机
      3. 在投影落点处采样图像特征
      4. 对所有有效采样求平均，得到该格子的特征
    """

    def __init__(self, embed_dims=256, num_heights=4,
                 z_range=(-2.0, 4.0)):
        super().__init__()
        self.embed_dims = embed_dims
        self.num_heights = num_heights
        self.z_range = z_range
        self.output_proj = nn.Linear(embed_dims, embed_dims)

    def forward(self, bev_query, img_feats, lidar2img, bev_coords):
        """
        bev_query  : [B, X*Y, C]
        img_feats  : [B, N, C, h, w]
        lidar2img  : [B, N, 4, 4]
        bev_coords : [X*Y, 2]   每格的 (x,y) 自车系坐标
        """
        B, N, C, h, w = img_feats.shape
        num_query = bev_coords.shape[0]
        Z = self.num_heights
        device = bev_query.device

        # ---- 1) 构造 3D 参考点（柱子上的 Z 个高度）----
        zs = torch.linspace(self.z_range[0], self.z_range[1], Z,
                            device=device)
        ref = torch.zeros(num_query, Z, 3, device=device)
        ref[:, :, 0] = bev_coords[:, 0:1]
        ref[:, :, 1] = bev_coords[:, 1:2]
        ref[:, :, 2] = zs.unsqueeze(0)

        ref_h = torch.cat(
            [ref.reshape(-1, 3),
             torch.ones(num_query * Z, 1, device=device)], dim=1)

        # ---- 2) 投影到各相机（一次 einsum 搞定 6 相机 × 所有点）----
        pts = torch.einsum('bnij,pj->bnpi', lidar2img, ref_h)

        depth = pts[..., 2]                          # [B,N,P]
        eps = 1e-5
        uv = pts[..., :2] / torch.clamp(depth.unsqueeze(-1), min=eps)

        # ---- 3) 有效性掩码：深度为正 且 落在图像内 ----
        img_w, img_h = w * 16, h * 16          # 骨干下采样 16 倍
        u_norm = uv[..., 0] / img_w * 2 - 1
        v_norm = uv[..., 1] / img_h * 2 - 1
        valid = ((depth > eps) &
                 (u_norm > -1) & (u_norm < 1) &
                 (v_norm > -1) & (v_norm < 1))

        grid = torch.stack([u_norm, v_norm], dim=-1)   # [B,N,P,2]

        # ---- 4) 采样图像特征 ----
        feats = img_feats.flatten(0, 1)               # [B*N,C,h,w]
        g = grid.flatten(0, 1).unsqueeze(2)           # [B*N,P,1,2]
        sampled = F.grid_sample(feats, g, mode='bilinear',
                                padding_mode='zeros',
                                align_corners=False)      # [B*N,C,P,1]
        sampled = sampled.squeeze(-1).permute(0, 2, 1)
        sampled = sampled.view(B, N, num_query, Z, C)

        # ---- 5) 按有效性加权平均 ----
        m = valid.view(B, N, num_query, Z, 1).float()
        summed = (sampled * m).sum(dim=(1, 3))          # [B,P,C]
        count = m.sum(dim=(1, 3)).clamp(min=1.0)
        out = summed / count

        # 残差连接（Transformer 标配，见论文①第 7 节）
        return self.output_proj(out) + bev_query


# ==================== BEV 编码器 ====================
class BEVEncoder(nn.Module):
    def __init__(self, embed_dims=256, num_layers=3, num_heights=4):
        super().__init__()
        self.layers = nn.ModuleList()
        for _ in range(num_layers):
            self.layers.append(nn.ModuleDict({
                'sca': SpatialCrossAttention(embed_dims, num_heights),
                'norm1': nn.LayerNorm(embed_dims),
                'ffn': nn.Sequential(
                    nn.Linear(embed_dims, embed_dims * 2),
                    nn.ReLU(inplace=True),
                    nn.Linear(embed_dims * 2, embed_dims)),
                'norm2': nn.LayerNorm(embed_dims),
            }))

    def forward(self, bev_query, img_feats, lidar2img, bev_coords):
        x = bev_query
        for layer in self.layers:
            x = layer['norm1'](
                layer['sca'](x, img_feats, lidar2img, bev_coords))
            x = layer['norm2'](x + layer['ffn'](x))
        return x


# ==================== 检测头 ====================
class CenterHead(nn.Module):
    """
    热力图检测头 —— 比 DETR 风格的 query 头更容易训练收敛

    输出:
        heatmap : [B, num_cls, Y, X]  每类的中心点热力图
        reg     : [B, 8, Y, X]  offset(2)+z(1)+logsize(3)+sincos(2)
    """

    def __init__(self, in_channels=256, num_classes=4):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Conv2d(in_channels, 128, 3, padding=1),
            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
        )
        self.hm_head = nn.Conv2d(128, num_classes, 1)
        self.reg_head = nn.Conv2d(128, 8, 1)
        # 热力图偏置初始化为负，避免训练初期大量假正例（CenterNet 惯例）
        self.hm_head.bias.data.fill_(-2.19)

    def forward(self, bev_feat):
        x = self.shared(bev_feat)
        return {'heatmap': self.hm_head(x), 'reg': self.reg_head(x)}


# ==================== 完整模型 ====================
class MiniBEVFormer(nn.Module):

    def __init__(self, bev_size=(128, 128), pc_range=51.2,
                 embed_dims=256, num_layers=3,
                 num_classes=4, pretrained=True):
        super().__init__()
        self.bev_h, self.bev_w = bev_size
        self.pc_range = pc_range
        self.embed_dims = embed_dims

        self.backbone = ImageBackbone(embed_dims, pretrained)
        # BEV query：可学习参数，这是 BEVFormer 的精髓（论文第 3 节）
        self.bev_embedding = nn.Embedding(
            self.bev_h * self.bev_w, embed_dims)
        self.encoder = BEVEncoder(embed_dims, num_layers)
        self.head = CenterHead(embed_dims, num_classes)

        # 预计算每格的自车系坐标（不参与训练）
        self.register_buffer('bev_coords', self._make_coords())

    def _make_coords(self):
        """生成 BEV 网格中心点的自车系坐标 [X*Y, 2]"""
        # x 向前（图像上方），y 向右
        xs = torch.linspace(self.pc_range, -self.pc_range, self.bev_h)
        ys = torch.linspace(-self.pc_range, self.pc_range, self.bev_w)
        gx, gy = torch.meshgrid(xs, ys, indexing='ij')
        return torch.stack([gx.flatten(), gy.flatten()], dim=1)

    def forward(self, imgs, lidar2img):
        B = imgs.shape[0]
        img_feats = self.backbone(imgs)

        bev_q = self.bev_embedding.weight.unsqueeze(0).repeat(B, 1, 1)
        bev = self.encoder(bev_q, img_feats, lidar2img, self.bev_coords)

        # [B, X*Y, C] → [B, C, Y, X]
        bev = bev.permute(0, 2, 1).reshape(
            B, self.embed_dims, self.bev_h, self.bev_w)
        return self.head(bev)


# ==================== 损失函数 ====================
def focal_loss(pred, target, alpha=2.0, beta=4.0):
    """CenterNet 风格的 focal loss（处理正负样本极度不平衡）"""
    pred = torch.clamp(torch.sigmoid(pred), 1e-4, 1 - 1e-4)
    pos_mask = target.eq(1).float()
    neg_mask = 1 - pos_mask

    pos_loss = -torch.log(pred) * torch.pow(1 - pred, alpha) * pos_mask
    neg_loss = (-torch.log(1 - pred) * torch.pow(pred, alpha)
                * torch.pow(1 - target, beta) * neg_mask)

    n_pos = pos_mask.sum()
    if n_pos == 0:
        return neg_loss.sum()
    return (pos_loss.sum() + neg_loss.sum()) / n_pos


# ==================== 自测 ====================
if __name__ == '__main__':
    print("=== MiniBEVFormer 形状与前向测试 ===")
    model = MiniBEVFormer(bev_size=(64, 64), num_layers=2,
                          pretrained=False)
    n = sum(p.numel() for p in model.parameters())
    print(f"参数量: {n/1e6:.1f} M")

    B, N, H, W = 2, 6, 256, 448
    imgs = torch.randn(B, N, 3, H, W)
    l2i = torch.eye(4).repeat(B, N, 1, 1)
    f = W / (2 * math.tan(math.radians(70) / 2))
    K = torch.tensor([[f, 0, W/2], [0, f, H/2], [0, 0, 1]])
    R = torch.tensor([[0., 1., 0.], [0., 0., -1.], [1., 0., 0.]])
    l2i[:, :, :3, :3] = K @ R

    model.eval()
    with torch.no_grad():
        out = model(imgs, l2i)
    print(f"heatmap: {tuple(out['heatmap'].shape)}")
    print(f"reg    : {tuple(out['reg'].shape)}")

    model.train()
    out = model(imgs, l2i)
    tgt = torch.zeros_like(out['heatmap']); tgt[0, 0, 30, 30] = 1.0
    loss = focal_loss(out['heatmap'], tgt)
    loss.backward()
    g = sum(p.grad.abs().sum().item() for p in model.parameters()
            if p.grad is not None)
    print(f"loss: {loss.item():.4f} | 梯度总量: {g:.1f} "
          f"{'✓' if g > 0 else '✗'}")

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""3D 框真值 → 热力图 + 回归目标"""

import numpy as np
import torch


def gaussian_2d(shape, sigma=1.0):
    m, n = [(s - 1.) / 2. for s in shape]
    y, x = np.ogrid[-m:m+1, -n:n+1]
    g = np.exp(-(x*x + y*y) / (2 * sigma * sigma))
    g[g < np.finfo(g.dtype).eps * g.max()] = 0
    return g


def draw_gaussian(heatmap, center, radius):
    """在热力图上画高斯核（CenterNet 标准做法）"""
    diameter = 2 * radius + 1
    g = gaussian_2d((diameter, diameter), sigma=diameter / 6)
    cx, cy = int(center[0]), int(center[1])
    H, W = heatmap.shape

    left, right = min(cx, radius), min(W - cx, radius + 1)
    top, bottom = min(cy, radius), min(H - cy, radius + 1)
    if left + right <= 0 or top + bottom <= 0:
        return

    masked_hm = heatmap[cy-top:cy+bottom, cx-left:cx+right]
    masked_g = g[radius-top:radius+bottom, radius-left:radius+right]
    # 取最大值：多个目标重叠时保留更强的响应
    np.maximum(masked_hm, masked_g, out=masked_hm)


def encode_targets(gt_boxes, gt_labels, bev_size=128,
                    pc_range=51.2, num_classes=4, min_radius=2):
    """
    gt_boxes  : [M, 7]  (x, y, z, l, w, h, yaw) 自车系
    返回:
        heatmap  [num_classes, H, W]
        reg      [8, H, W]
        reg_mask [H, W]     哪些位置有真值（只在这些位置算回归损失）
    """
    H = W = bev_size
    heatmap = np.zeros((num_classes, H, W), dtype=np.float32)
    reg = np.zeros((8, H, W), dtype=np.float32)
    mask = np.zeros((H, W), dtype=np.float32)

    # 自车系 → BEV 像素坐标
    # 注意与模型 _make_coords 的定义保持一致：
    #   x 向前 → 行索引从下往上；y 向右 → 列索引从左往右
    voxel = 2 * pc_range / bev_size

    for box, label in zip(gt_boxes, gt_labels):
        x, y, z, l, w, h, yaw = box
        if abs(x) >= pc_range or abs(y) >= pc_range:
            continue

        col_f = (y + pc_range) / voxel
        row_f = (pc_range - x) / voxel
        col, row = int(col_f), int(row_f)
        if not (0 <= col < W and 0 <= row < H):
            continue

        # 高斯半径与物体尺寸相关（大物体半径大）
        radius = max(min_radius, int(max(l, w) / voxel / 2))
        draw_gaussian(heatmap[int(label)], (col, row), radius)

        # 回归目标（在中心格子上）
        reg[0, row, col] = col_f - col          # 亚像素 offset x
        reg[1, row, col] = row_f - row          # 亚像素 offset y
        reg[2, row, col] = z
        reg[3, row, col] = np.log(max(l, 0.1))    # log 尺寸更稳定
        reg[4, row, col] = np.log(max(w, 0.1))
        reg[5, row, col] = np.log(max(h, 0.1))
        reg[6, row, col] = np.sin(yaw)          # 用 sin/cos 避免角度突变
        reg[7, row, col] = np.cos(yaw)
        mask[row, col] = 1

    return (torch.from_numpy(heatmap),
            torch.from_numpy(reg),
            torch.from_numpy(mask))

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""MiniBEVFormer 训练脚本 —— 适配共享 GPU 环境"""

import argparse
import os
import time

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from bev_dataset import BEVDataset, bev_collate_fn
from bev_target import encode_targets
from mini_bevformer import MiniBEVFormer, focal_loss
from gpu_polite import setup_polite_gpu, report_memory


def build_batch_targets(batch, bev_size, pc_range, num_classes, device):
    """把一批变长真值编码成张量"""
    hms, regs, masks = [], [], []
    for boxes, labels in zip(batch['gt_boxes'], batch['gt_labels']):
        hm, reg, m = encode_targets(
            boxes.numpy(), labels.numpy(),
            bev_size=bev_size, pc_range=pc_range,
            num_classes=num_classes)
        hms.append(hm); regs.append(reg); masks.append(m)
    return (torch.stack(hms).to(device),
            torch.stack(regs).to(device),
            torch.stack(masks).to(device))


def train_one_epoch(model, loader, optimizer, scaler, device,
                    args, writer, epoch):
    model.train()
    total, n_batch = 0.0, 0
    pbar = tqdm(loader, desc=f"Epoch {epoch}")

    for it, batch in enumerate(pbar):
        imgs = batch['imgs'].to(device, non_blocking=True)
        l2i = batch['lidar2img'].to(device, non_blocking=True)
        gt_hm, gt_reg, gt_mask = build_batch_targets(
            batch, args.bev_size, args.pc_range, args.num_classes, device)

        # 混合精度：L20 支持 bf16，比 fp16 更稳定（见 1.1 节）
        with torch.autocast('cuda', dtype=torch.bfloat16,
                            enabled=args.amp):
            out = model(imgs, l2i)
            loss_hm = focal_loss(out['heatmap'].float(), gt_hm)

            # 回归损失只在有真值的位置计算
            m = gt_mask.unsqueeze(1)                # [B,1,H,W]
            n_pos = m.sum().clamp(min=1.0)
            loss_reg = (F.l1_loss(out['reg'].float() * m,
                                  gt_reg * m, reduction='sum') / n_pos)

            loss = loss_hm + args.reg_weight * loss_reg

        optimizer.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        # 梯度裁剪：BEV 模型训练初期梯度容易爆
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 35.0)
        scaler.step(optimizer)
        scaler.update()

        total += loss.item(); n_batch += 1
        pbar.set_postfix({
            'loss': f"{loss.item():.3f}",
            'hm': f"{loss_hm.item():.3f}",
            'reg': f"{loss_reg.item():.3f}",
        })

        if writer and it % 20 == 0:
            step = epoch * len(loader) + it
            writer.add_scalar('train/loss', loss.item(), step)
            writer.add_scalar('train/loss_heatmap', loss_hm.item(), step)
            writer.add_scalar('train/loss_reg', loss_reg.item(), step)

    return total / max(n_batch, 1)


@torch.no_grad()
def validate(model, loader, device, args):
    model.eval()
    total, n = 0.0, 0
    for batch in tqdm(loader, desc="Val", leave=False):
        imgs = batch['imgs'].to(device)
        l2i = batch['lidar2img'].to(device)
        gt_hm, gt_reg, gt_mask = build_batch_targets(
            batch, args.bev_size, args.pc_range, args.num_classes, device)
        with torch.autocast('cuda', dtype=torch.bfloat16, enabled=args.amp):
            out = model(imgs, l2i)
            l_hm = focal_loss(out['heatmap'].float(), gt_hm)
            m = gt_mask.unsqueeze(1)
            l_reg = (F.l1_loss(out['reg'].float() * m, gt_reg * m,
                               reduction='sum') / m.sum().clamp(min=1.0))
        total += (l_hm + args.reg_weight * l_reg).item(); n += 1
    return total / max(n, 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--root', default='./dataset')
    ap.add_argument('--out', default='./work_dir')
    ap.add_argument('--epochs', type=int, default=24)
    ap.add_argument('--batch-size', type=int, default=4)
    ap.add_argument('--lr', type=float, default=2e-4)
    ap.add_argument('--bev-size', type=int, default=128)
    ap.add_argument('--num-layers', type=int, default=3)
    ap.add_argument('--pc-range', type=float, default=51.2)
    ap.add_argument('--num-classes', type=int, default=4)
    ap.add_argument('--img-w', type=int, default=448)
    ap.add_argument('--img-h', type=int, default=256)
    ap.add_argument('--reg-weight', type=float, default=0.25)
    ap.add_argument('--workers', type=int, default=4)
    ap.add_argument('--amp', action='store_true', default=True)
    ap.add_argument('--max-gpu-gb', type=float, default=30,
                    help='共享 GPU 上限制自己的显存占用')
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    device = torch.device('cuda')

    # 共享 GPU 礼貌用法（1.1 节）
    setup_polite_gpu(max_gb=args.max_gpu_gb)

    # --- 数据 ---
    train_ds = BEVDataset(args.root, 'train',
                          img_size=(args.img_w, args.img_h),
                          pc_range=args.pc_range, augment=True)
    val_ds = BEVDataset(args.root, 'val',
                        img_size=(args.img_w, args.img_h),
                        pc_range=args.pc_range, augment=False)
    train_dl = DataLoader(train_ds, batch_size=args.batch_size,
                          shuffle=True, num_workers=args.workers,
                          collate_fn=bev_collate_fn, pin_memory=True,
                          drop_last=True)
    val_dl = DataLoader(val_ds, batch_size=args.batch_size,
                        shuffle=False, num_workers=2,
                        collate_fn=bev_collate_fn)

    # --- 模型 ---
    model = MiniBEVFormer(
        bev_size=(args.bev_size, args.bev_size),
        pc_range=args.pc_range, num_layers=args.num_layers,
        num_classes=args.num_classes).to(device)
    n_param = sum(p.numel() for p in model.parameters())
    print(f"模型参数量: {n_param/1e6:.1f} M")

    # 骨干用更小学习率（预训练权重，不宜大改）
    backbone_params = list(model.backbone.parameters())
    other_params = [p for n, p in model.named_parameters()
                    if not n.startswith('backbone')]
    optimizer = torch.optim.AdamW([
        {'params': backbone_params, 'lr': args.lr * 0.1},
        {'params': other_params, 'lr': args.lr},
    ], weight_decay=0.01)

    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=[args.lr * 0.1, args.lr],
        total_steps=args.epochs * len(train_dl), pct_start=0.3)
    scaler = torch.amp.GradScaler('cuda', enabled=args.amp)
    writer = SummaryWriter(os.path.join(args.out, 'tb'))

    best = float('inf')
    for epoch in range(args.epochs):
        t0 = time.time()
        tr = train_one_epoch(model, train_dl, optimizer, scaler,
                             device, args, writer, epoch)
        scheduler.step()
        va = validate(model, val_dl, device, args)
        dt = time.time() - t0

        print(f"Epoch {epoch:3d} | train {tr:.4f} | val {va:.4f} "
              f"| {dt/60:.1f} min")
        report_memory(f" ep{epoch}")
        writer.add_scalar('epoch/train_loss', tr, epoch)
        writer.add_scalar('epoch/val_loss', va, epoch)

        ckpt = {'epoch': epoch, 'model': model.state_dict(),
                'args': vars(args), 'val_loss': va}
        torch.save(ckpt, os.path.join(args.out, 'last.pth'))
        if va < best:
            best = va
            torch.save(ckpt, os.path.join(args.out, 'best.pth'))
            print(f"  → 保存最佳模型 (val {va:.4f})")

    writer.close()
    print(f"\n训练完成，最佳 val loss: {best:.4f}")


if __name__ == '__main__':
    main()

conda activate ad

# ① 先用极小配置验证代码能跑（2 分钟内应看到 loss）
python train_bev.py --root ./dataset \
    --epochs 1 --batch-size 2 --bev-size 64 --num-layers 1 \
    --img-w 320 --img-h 192 --out ./work_dir_smoke

# ② 确认无误后正式训练（用 tmux 挂后台）
tmux new -s train
python train_bev.py --root ./dataset \
    --epochs 24 --batch-size 4 --bev-size 128 --num-layers 3 \
    --max-gpu-gb 30 --out ./work_dir
# Ctrl+B D 脱离

# ③ 看训练曲线（1.1 节配了端口转发，本地浏览器可直接开）
tensorboard --logdir ./work_dir/tb --port 6006
# 本地打开 http://localhost:6006

# ④ 监控显存与进度
watch -n 5 nvidia-smi
tmux attach -t train