PointNet

f(x₁, x₂) = f(x₂, x₁)

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5
6class Tnet(nn.Module):
7    """
8    PointNet (2017) 中的 Tnet 模块
9    """
10    def __init__(
11        self,
12        k: int = 64
13    ):
14        """
15        参数：
16            k: 输入通道数（同时也是输出变换矩阵的行数/列数）
17        """
18        super().__init__()
19        self.k = k
20
21        # 卷积层堆叠
22        self.conv1 = nn.Conv1d(k, 64, 1)
23        self.bn1 = nn.BatchNorm1d(64)
24        self.conv2 = nn.Conv1d(64, 128, 1)
25        self.bn2 = nn.BatchNorm1d(128)
26        self.conv3 = nn.Conv1d(128, 1024, 1)
27        self.bn3 = nn.BatchNorm1d(1024)
28
29        # 全连接层堆叠
30        self.fc1 = nn.Linear(1024, 512)
31        self.bn4 = nn.BatchNorm1d(512)
32        self.fc2 = nn.Linear(512, 256)
33        self.bn5 = nn.BatchNorm1d(256)
34        self.fc3 = nn.Linear(256, k * k)
35
36    def forward(
37        self,
38        x: torch.Tensor
39    ) -> torch.Tensor:
40        """
41        模型前向传播
42
43        参数：
44            x: 输入张量
45        """
46        batch_size = x.size()[0]
47
48        # 卷积层堆叠
49        x = F.relu(self.bn1(self.conv1(x)))
50        x = F.relu(self.bn2(self.conv2(x)))
51        x = F.relu(self.bn3(self.conv3(x)))
52        x = torch.max(x, 2, keepdim=True)[0]
53    
54        # 全连接层堆叠
55        x = x.view(-1, 1024)
56        x = F.relu(self.bn4(self.fc1(x)))
57        x = F.relu(self.bn5(self.fc2(x)))
58        x = self.fc3(x)
59    
60        # 加上单位矩阵 I，以构成有效的变换矩阵
61        eye = torch.eye(self.k).view(1, self.k * self.k).repeat(batch_size, 1)
62        x = x + eye
63        x = x.view(-1, self.k, self.k)
64        return x
65
66
67# 创建一个 Tnet 实例
68tnet = Tnet(k=3)
69
70# 加载数据
71from pytorch3d.io import load_objs_as_meshes
72from pytorch3d.ops import sample_points_from_meshes
73
74# 从文件加载网格
75mesh_paths = [
76    './data/sphere1.obj',
77    './data/cube1.obj'
78]
79meshes = load_objs_as_meshes(mesh_paths, load_textures=False)
80
81# 从网格中采样并重塑点云
82n_vert_samples = 128
83point_samples = sample_points_from_meshes(meshes, n_vert_samples)
84point_samples = point_samples.transpose(2, 1)
85
86# 执行 tnet 的前向传播
87output = tnet(point_samples)
88print(output.shape)
89print(output[0])

1from pytorch3d.io import load_objs_as_meshes
2from pytorch3d.ops import sample_points_from_meshes
3
4from plot import scatter3d, plt
5
6
7# 从文件加载网格
8mesh_paths = [
9    './data/sphere1.obj',
10    './data/cube1.obj'
11]
12meshes = load_objs_as_meshes(mesh_paths, load_textures=False)
13
14# 从网格中采样点
15n_vert_samples = 128
16point_samples = sample_points_from_meshes(meshes, n_vert_samples)
17print(point_samples.shape)
18
19# 将网格顶点以散点图形式显示
20scatter3d(point_samples[0, :, 0].cpu(), point_samples[0, :, 1].cpu(),
21          point_samples[0, :, 2].cpu(), title='球体')
22plt.savefig('./output/sphere.png')
23scatter3d(point_samples[1, :, 0].cpu(), point_samples[1, :, 1].cpu(),
24          point_samples[1, :, 2].cpu(), title='立方体')
25plt.savefig('./output/cube.png')

1import os
2
3import torch
4import torch.nn.functional as F
5from pytorch3d.structures import Meshes
6from pytorch3d.ops import sample_points_from_meshes
7from pytorch3d.io import load_objs_as_meshes
8
9from pointnet import PointNetClassifier, feature_transform_regularizer, device
10from plot import scatter3d, plot_losses, plt
11
12
13def sample_batch(
14    meshes: Meshes,
15    batch_size: int = 32,
16    n_vert_samples: int = 128
17) -> tuple[torch.Tensor, torch.Tensor]:
18    """
19    从 Meshes 实例中创建一批随机表面采样
20
21    参数：
22       meshes: 输入网格
23       batch_size: 批次大小
24       n_vert_samples: 批次中每个样本的顶点采样数量
25    """
26    rand_idxs = torch.randint(0, len(meshes), [batch_size], device=device)
27    list_samples = [sample_points_from_meshes(meshes[i.item()], n_vert_samples) for i in rand_idxs]
28    x_batch = torch.cat(list_samples, axis=0)
29    return x_batch, rand_idxs
30
31
32# 设置参数
33n_vert_samples = 128
34mesh_paths = [
35    './data/sphere1.obj',
36    './data/cube1.obj'
37]
38batch_size = 16
39n_batches = 50
40use_feature_transform = True
41print_period = 5
42output_dir = './output'
43
44# 创建输入网格
45meshes = load_objs_as_meshes(mesh_paths, load_textures=False)
46
47# 显示输入网格
48points_sphere = meshes[0].verts_packed()
49scatter3d(points_sphere[:, 0].cpu(), points_sphere[:, 1].cpu(),
50          points_sphere[:, 2].cpu(), title='球体')
51plt.savefig('./output/show-sphere.png')
52points_cube = meshes[1].verts_packed()
53scatter3d(points_cube[:, 0].cpu(), points_cube[:, 1].cpu(),
54          points_cube[:, 2].cpu(), title='立方体')
55plt.savefig('./output/show-cube.png')
56
57# 创建 PointNet 模型
58model = PointNetClassifier(k=len(meshes), feature_transform=use_feature_transform)
59model = model.to(device)
60
61# 创建优化器和学习率调度器
62optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))
63
64# 开始训练
65losses = []
66test_losses = []
67for i in range(n_batches):
68    optimizer.zero_grad()
69
70    # 采样一个批次
71    points, target = sample_batch(meshes, batch_size, n_vert_samples)
72    points = points.transpose(2, 1)
73
74    # 前向传播
75    model = model.train()
76    pred, trans, trans_feat = model(points)
77
78    # 计算损失
79    loss = F.nll_loss(pred, target)
80    if use_feature_transform:
81        loss += feature_transform_regularizer(trans_feat) * 0.001
82
83    losses.append(loss.detach().cpu().item())
84
85    # 反向传播并更新参数
86    loss.backward()
87    optimizer.step()
88
89    # 计算测试损失
90    points, target = sample_batch(meshes, batch_size, n_vert_samples)
91    points = points.transpose(2, 1)
92    model = model.eval()
93    pred, _, _ = model(points)
94    loss = F.nll_loss(pred, target)
95    test_losses.append(loss.detach().cpu().item())
96
97    # 输出训练进度
98    if i % print_period == 0:
99        pred_choice = pred.data.max(1)[1]
100        n_correct = pred_choice.eq(target.data).cpu().sum()
101        acc = n_correct.item() / float(batch_size)
102        print(f'批次 {i:4d}/{n_batches:4d}: 验证损失: {loss.item():1.4e} 准确率: {acc:.2%}')
103
104# 绘制损失曲线
105plot_losses(losses, test_losses)
106plt.savefig('./output/plot_losses.png')
107
108# 计算测试准确率
109model = model.eval()
110points, target = sample_batch(meshes, 64, n_vert_samples)
111points = points.transpose(2, 1)
112pred, _, _ = model(points)
113pred_choice = pred.data.max(1)[1]
114n_correct = pred_choice.eq(target.data).cpu().sum().item()
115n_points = points.size()[0]
116print(f'测试准确率: {n_correct / float(n_points):.2%} ({n_correct:d}/{n_points:d})')