从零实现WiFi姿态估计：Python+PyTorch实战CSI-Former模型

在计算机视觉领域，基于WiFi信号的人体姿态估计正成为新兴研究方向。与依赖摄像头的传统方案相比，WiFi信号具有穿透性强、隐私友好和设备普及等优势。本文将手把手带你用PyTorch实现CSI-Former模型，从数据集处理到模型训练完整复现论文核心成果。

1. 环境配置与数据准备

首先需要搭建支持CUDA的PyTorch环境。建议使用conda创建独立环境：

conda create -n wifi_pose python=3.8
conda activate wifi_pose
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install h5py scipy matplotlib tqdm

Wi-Pose数据集包含CSI信号和对应的骨骼关键点标注，数据结构如下：

文件类型	内容描述	数据维度
.mat	原始CSI信号	[300, 90]
.mat	18个骨骼点坐标(x,y,置信度)	[3, 18]

下载数据集后，建议按以下结构组织：

Wi-Pose/
├── Train/
│   ├── walking/
│   ├── running/
│   └── ...
└── Test/
    ├── jumping/
    ├── sitting/
    └── ...

2. 数据预处理实战

WiFi信号易受多径效应干扰，需进行降噪处理。我们实现巴特沃斯滤波器和基于注意力的联合降噪方案：

import scipy.signal as signal

def butterworth_filter(csi_data, cutoff=0.1):
    b, a = signal.butter(4, cutoff, 'lowpass')
    return signal.filtfilt(b, a, csi_data, axis=0)

class AttentionDenoiser(nn.Module):
    def __init__(self, subcarrier_num=90):
        super().__init__()
        self.attention = nn.Sequential(
            nn.Linear(subcarrier_num, 64),
            nn.ReLU(),
            nn.Linear(64, subcarrier_num),
            nn.Softmax(dim=-1)
        )
    
    def forward(self, x):
        # x shape: [batch, time, subcarrier]
        weights = self.attention(x.mean(1))
        return x * weights.unsqueeze(1)

数据加载时需注意h5py版本兼容性问题：

def load_mat(file_path):
    try:
        with h5py.File(file_path, 'r') as f:
            csi = np.array(f['CSI']).T  # 转置匹配PyTorch维度
            kps = np.array(f['SkeletonPoints'])
    except:
        data = scipy.io.loadmat(file_path)
        csi = data['CSI']
        kps = data['SkeletonPoints']
    return torch.FloatTensor(csi), torch.FloatTensor(kps)

3. CSI-Former模型架构解析

CSI-Former的核心创新在于Performer层，这是Transformer的改进版本，具有线性计算复杂度：

class PerformerLayer(nn.Module):
    def __init__(self, dim, heads=8):
        super().__init__()
        self.dim = dim
        self.heads = heads
        self.to_qkv = nn.Linear(dim, dim*3)
        self.gelu = nn.GELU()
        
        # 随机正交矩阵用于特征映射
        self.proj = nn.Parameter(torch.randn(dim, dim))
        nn.init.orthogonal_(self.proj)
        
    def forward(self, x):
        B, T, C = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        
        # 使用随机特征近似注意力
        phi = lambda x: self.gelu(x @ self.proj)
        q, k = map(phi, qkv[:2])
        v = qkv[2]
        
        attn = torch.einsum('bhd,bhl->bhl', q, k) / (C ** 0.5)
        out = torch.einsum('bhl,bhd->bhd', attn.softmax(dim=-1), v)
        return out + x  # 残差连接

完整的CSI-Former模型包含三个关键组件：

1. CSI编码器：3层CNN提取局部特征
2. Performer堆叠：12层处理时序依赖
3. 姿态解码器：反卷积网络输出关键点

class CSIFormer(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv1d(90, 256, 5, padding=2),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.MaxPool1d(2)
        )
        
        self.performer = nn.Sequential(*[
            PerformerLayer(256) for _ in range(12)
        ])
        
        self.decoder = nn.Sequential(
            nn.ConvTranspose1d(256, 128, 5),
            nn.ReLU(),
            nn.ConvTranspose1d(128, 54, 5)  # 18点×3坐标
        )
    
    def forward(self, x):
        x = self.encoder(x.transpose(1,2))
        x = self.performer(x.transpose(1,2))
        return self.decoder(x.transpose(1,2)).view(-1,18,3)

4. 训练策略与评估指标

采用教师-学生网络框架，使用AlphaPose生成的标注作为监督信号。关键训练技巧包括：

• 骨架邻接矩阵(SAM)正则化：防止关节点预测相互矛盾
• 动态学习率：初始lr=3e-4，每5epoch衰减10%
• 混合精度训练：减少显存占用

def sam_loss(pred, target):
    # 计算关节点间的相对位置约束
    pred_adj = pred.unsqueeze(3) - pred.unsqueeze(2)  # [B,18,18,3]
    target_adj = target.unsqueeze(3) - target.unsqueeze(2)
    return F.mse_loss(pred_adj, target_adj)

def train_epoch(model, loader, optimizer):
    model.train()
    total_loss = 0
    for csi, kps in loader:
        optimizer.zero_grad()
        with torch.cuda.amp.autocast():
            pred = model(csi.cuda())
            loss = F.mse_loss(pred[...,:2], kps[...,:2].cuda())  # 仅计算xy坐标
            loss += 0.1 * sam_loss(pred, kps.cuda())
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

评估采用PCK(Percentage of Correct Keypoints)指标：

def calculate_pck(pred, target, threshold=0.05):
    # 归一化到[0,1]范围
    pred_norm = (pred - pred.min()) / (pred.max() - pred.min())
    target_norm = (target - target.min()) / (target.max() - target.min())
    
    distances = torch.norm(pred_norm[...,:2] - target_norm[...,:2], dim=-1)
    return (distances < threshold).float().mean()

5. 常见问题排查

在实际复现过程中可能会遇到以下典型问题：

维度不匹配错误

• 现象：RuntimeError: shape mismatch
• 解决方案：检查.mat文件读取时的转置操作，确保CSI信号始终是[序列长度, 子载波数]格式

h5py版本问题

• 现象：Unable to open file (truncated file)
• 解决方案：降级h5py到2.10.0版本或使用scipy.io.loadmat替代

梯度爆炸

• 现象：loss变为nan
• 解决方案：添加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

显存不足

• 现象：CUDA out of memory
• 解决方案：
  1. 减小batch_size到8或16
  2. 使用torch.utils.checkpoint分段计算

经过完整训练后，在测试集上应能达到PCK@0.05约78%的准确率，比传统ResNet基线提升3-5个百分点。实际部署时建议对CSI信号进行实时滤波处理，并添加滑动窗口平滑预测结果。