1. 基于PyTorch的Occupancy网络基础框架搭建

任务说明

本案例实现一个最小可运行的Occupancy网络:输入部分点云(partial point cloud)与任意查询点坐标,模型输出该查询点被占用的概率。训练数据使用程序生成的合成3D形状(球体表面采样),通过对比查询点是否位于形状内部来构造监督信号。

环境依赖

Python 3.8+,PyTorch 1.10+,NumPy。无需其他外部库。

实现代码

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# ---------- 1. 合成数据集:球体点云 + 占用查询 ----------
def sample_sphere_surface(n_points, radius=1.0):
    """在球面上均匀采样点云"""
    u = np.random.rand(n_points)
    v = np.random.rand(n_points)
    theta = 2 * np.pi * u
    phi = np.arccos(2 * v - 1)
    x = radius * np.sin(phi) * np.cos(theta)
    y = radius * np.sin(phi) * np.sin(theta)
    z = radius * np.cos(phi)
    return np.stack([x, y, z], axis=1).astype(np.float32)

def generate_training_batch(batch_size, n_surface=512, n_query=256):
    """生成一个batch:每个样本包含点云、查询点、占用标签"""
    point_clouds = []
    queries = []
    labels = []
    for _ in range(batch_size):
        radius = np.random.uniform(0.8, 1.2)
        pc = sample_sphere_surface(n_surface, radius)
        # 查询点:一半在球内(占用=1),一半在球外(占用=0)
        q = np.random.uniform(-1.5, 1.5, size=(n_query, 3)).astype(np.float32)
        dist = np.linalg.norm(q, axis=1)
        lab = (dist < radius).astype(np.float32)
        point_clouds.append(pc)
        queries.append(q)
        labels.append(lab)
    return (torch.from_numpy(np.stack(point_clouds)),
            torch.from_numpy(np.stack(queries)),
            torch.from_numpy(np.stack(labels)))

# ---------- 2. Occupancy Network 模型 ----------
class PointNetEncoder(nn.Module):
    """简单PointNet编码器:将点云编码为全局特征向量"""
    def __init__(self, feat_dim=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(3, 64), nn.ReLU(),
            nn.Linear(64, 128), nn.ReLU(),
            nn.Linear(128, feat_dim)
        )
    def forward(self, pc):
        # pc: (B, N, 3)
        x = self.mlp(pc)          # (B, N, feat_dim)
        return x.max(dim=1).values  # (B, feat_dim)

class OccupancyNetwork(nn.Module):
    def __init__(self, feat_dim=256):
        super().__init__()
        self.encoder = PointNetEncoder(feat_dim)
        self.decoder = nn.Sequential(
            nn.Linear(feat_dim + 3, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(),
            nn.Linear(128, 1), nn.Sigmoid()
        )
    def forward(self, pc, q):
        # pc: (B, N, 3), q: (B, M, 3)
        feat = self.encoder(pc)          # (B, feat_dim)
        B, M = q.shape[0], q.shape[1]
        feat_exp = feat.unsqueeze(1).expand(B, M, -1)  # (B, M, feat_dim)
        x = torch.cat([feat_exp, q], dim=2)            # (B, M, feat_dim+3)
        return self.decoder(x).squeeze(-1)             # (B, M)

# ---------- 3. 训练与验证 ----------
def train():
    torch.manual_seed(42)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = OccupancyNetwork().to(device)
    opt = optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.BCELoss()

    model.train()
    for epoch in range(30):
        pc, q, lab = generate_training_batch(batch_size=8)
        pc, q, lab = pc.to(device), q.to(device), lab.to(device)
        opt.zero_grad()
        pred = model(pc, q)
        loss = loss_fn(pred, lab)
        loss.backward()
        opt.step()
        if (epoch + 1) % 10 == 0:
            print(f'Epoch {epoch+1:3d} | Loss: {loss.item():.4f}')

    # 验证:单一样本上预测准确率
    model.eval()
    with torch.no_grad():
        pc, q, lab = generate_training_batch(batch_size=1)
        pc, q, lab = pc.to(device), q.to(device), lab.to(device)
        pred = model(pc, q)
        acc = ((pred > 0.5).float() == lab).float().mean().item()
        print(f'验证准确率: {acc*100:.1f}%')

if __name__ == '__main__':
    train()

运行方式

将上述代码保存为 occupancy_network.py,在终端执行:

python occupancy_network.py

验证方法与通过标准

  • 训练过程中每10个epoch打印的Loss应持续下降,最终低于0.1。
  • 训练结束后打印的验证准确率应高于95%。该准确率表示模型对球内/球外查询点的二分类正确率,反映模型已学会利用点云推断形状占用。
  • 若准确率低于90%,可增加训练轮次至50或增大batch_size至16。

关键说明

本实现使用PointNet最大池化提取点云全局特征,将查询点坐标与全局特征拼接后经MLP输出占用概率。合成数据中球体半径在0.8~1.2间随机变化,迫使模型不能仅依赖坐标值判断,而必须结合点云信息。该最小实现完整覆盖了Occupancy Network的核心流程:点云编码、查询点条件解码、占用监督训练。