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的核心流程:点云编码、查询点条件解码、占用监督训练。