一、数据准备
该部分代码位于./data/modelnet_shrec_loader.py
中。读取的数据为pc_np
(点的坐标),surface_normal_np
(法向量),som_node_np
(som节点坐标)和class_id
(类别)。然后对数据增强,包括旋转、微扰、尺度变换和位移。返回点的坐标、法向量、类别、som节点和每个som节点在som节点中的k个近邻点的索引。
二、模型
该部分代码位于./models/classifier.py
。主要包括两个部分:编码网络和分类网络。
2.1编码网络
编码器如上图蓝色底色的部分所示。具体的定义在./models/networks.py
中。在Encoder类
的forward
中可以看到,其输入为点的坐标、法向量、节点坐标和节点到节点的knn索引。
2.1.1 SOM层
如上图所示,第一部分为SOM层。通过./util/som.py
的BatchSOM
类中的query_topk()
函数实现映射,返回mask,mask_row_max和min_idx。
(1)mask:其大小为[B,kN,M],第i个N*M行表示节点是否为N个点的第i近邻节点。
上图中红色的1表示第2个节点是第一个点的第一近邻点。蓝色的1表示第3个节点是第四个点的第一近邻点,而最下方的1表示第5个节点是第三个点的第二近邻点。
(2)mask_row_max:其大小为[B,M],表示每个节点是否存在近邻的点
其中1表示第i个节点是某个点的近邻点。
(3)min_idx:其大小为[B,kN],第i个N*1行表示N个点的第i近邻节点的索引。
与mask对应,表示最近邻索引值。
然后获取以每个节点为近邻的所有点的中心。
self.mask, mask_row_max, min_idx = self.som_builder.query_topk(x.data, k=self.opt.k) # BxkNxnode_num, Bxnode_num
mask_row_sum = torch.sum(self.mask, dim=1) # BxM
mask = self.mask.unsqueeze(1) # Bx1xkNxM
#将x和sn堆叠
x_list, sn_list = [], []
for i in range(self.opt.k):
x_list.append(x)
sn_list.append(sn)
x_stack = torch.cat(tuple(x_list), dim=2) # B x C x kN
sn_stack =