论文代码-PointNet && PointNet++

PointNet&PointNet++代码解读

引用说明

文章中的代码全部来自于Github仓库: Pointnet_Pointnet2_pytorch

本文更关注语义分割以及零件分割部分代码

PointNet

• 完整的PointNet结构代码
  来自于models文件夹下的pointnet_util.py文件

class PointNetEncoder(nn.Module):
    def __init__(self, global_feat=True, feature_transform=False, channel=3):
        super(PointNetEncoder, self).__init__()
        self.stn = STN3d(channel)  # 这个是3*3的T-net
        self.conv1 = torch.nn.Conv1d(channel, 64, 1)  # 从输入channel维度变成64维
        self.conv2 = torch.nn.Conv1d(64, 128, 1)      # 到128
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)    # 到1024
        # 三个批量归一化
        self.bn1 = nn.BatchNorm1d(64)               
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.global_feat = global_feat  # 是否要全局特征 boolean型变量 用户输入数据
        self.feature_transform = feature_transform  # 是否要进行特征那个矩阵的转换 如果要则需要k*k的T-net
        if self.feature_transform:
            self.fstn = STNkd(k=64)

    def forward(self, x):
        B, D, N = x.size()  # 批量大小 点云的维度3或6 一个点云的点数
        trans = self.stn(x)  # 先经过T-net调整
        x = x.transpose(2, 1)  # 这里将输入的维度交换了一下 每一列三个点坐标然后每一行是点的个数
        # 提feature和xyz nx ny nz
        if D > 3:
            feature = x[:, :, 3:]
            x = x[:, :, :3]
        # 用T-net来实现变换的不变性
        x = torch.bmm(x, trans)  # 两个三维张量的相乘 batch_size相同 (b,i,j) * (b,j,k) -> (b,i,k) 就是多了一个b
        # 将特征拼接回去
        if D > 3:
            x = torch.cat([x, feature], dim=2)
        x = x.transpose(2, 1)  # 还原原来的维度
        x = F.relu(self.bn1(self.conv1(x)))  # 卷积 + 归一化 + 激活函数 MLP 变成64维

        # 如果需要特征矩阵的还原 就再做一次T-net
        if self.feature_transform:
            trans_feat = self.fstn(x)
            x = x.transpose(2, 1)
            x = torch.bmm(x, trans_feat)
            x = x.transpose(2, 1)
        else:
            trans_feat = None

        pointfeat = x
        x = F.relu(self.bn2(self.conv2(x)))  # 64 -> 128
        x = self.bn3(self.conv3(x))  # 128 -> 1024
        x = torch.max(x, 2, keepdim=True)[0]  # 用对称函数来提取全局特征
        x = x.view(-1, 1024)  # 全局特征维度整理 1024
        # 需要局部特征就返回 否则直接concat返回
        if self.global_feat:
            return x, trans, trans_feat
        else:
            x = x.view(-1, 1024, 1).repeat(1, 1, N)
            return torch.cat([x, pointfeat], 1), trans, trans_feat

• 3*3的T-net代码
  来自于models文件夹下的pointnet_util.py文件

'''
3*3的T-net 类似于一个mini-PointNet结构 
这个T-net的参数不是提前设定好的 而是跟着整个网络的运行不断计算的
'''
class STN3d(nn.Module):
    def __init__(self, channel):
        super(STN3d, self).__init__()
        # PointNet是使用一维卷积进行高维度的映射操作，多输入多输出通道，相当于每个通道学习一个卷积核，这里通道就是维度
        # 采用卷积而不是全连接实现可能是因为cudnn计算上有优化
        self.conv1 = torch.nn.Conv1d(channel, 64, 1)  # 从输入channel数量3或者是6映射到64维
        self.conv2 = torch.nn.Conv1d(64, 128, 1)      # 64->128
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)    # 128->1024
        self.fc1 = nn.Linear(1024, 512)               # full-connection全连接层 直接用线性层实现
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 9)
        self.relu = nn.ReLU()                         # 激活函数定义
        # 下面是batch_norm 批量归一化操作 卷积神经网络常用于加快模型收敛速度 增强泛化能力
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)

    # 李沐导师讲过 定义一个网络就是继承nn.Module然后重写init和forward函数
    def forward(self, x):
        batchsize = x.size()[0]  # 提取出批量大小
        # 卷积+归一化+激活函数
        x = F.relu(self.bn1(self.conv1(x)))  # 首先映射到64维
        x = F.relu(self.bn2(self.conv2(x)))  # 到128
        x = F.relu(self.bn3(self.conv3(x)))  # 到1024
        x = torch.max(x, 2, keepdim=True)[0]  # 使用对称函数max 最大池化操作
        x = x.view(-1, 1024)                 # 展开成1024维向量 这里就是全局特征了
        # 通过全连接+批量归一化+激活函数实现MLP
        x = F.relu(self.bn4(self.fc1(x)))  # 1024 -> 512
        x = F.relu(self.bn5(self.fc2(x)))  # 512 -> 256
        x = self.fc3(x)  # 256 -> 9 到这里获得了对于变换的不变性的旋转矩阵3*3的9个元素

        # 这里找到一个对角阵然后展平
        iden = Variable(torch.from_numpy(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32))).view(1, 9).repeat(
            batchsize, 1)
        if x.is_cuda:
            iden = iden.cuda()
        x = x + iden  # 线性变换 + 平移  这里没有读原文的T-net实现，不知道为啥要做这个
        x = x.view(-1, 3, 3)  # 整理成3*3的矩阵
        return x

• 补充
  • feature_transform_reguliarzer方法主要是对于T-net生成的矩阵正则的处理，希望是更接近正交矩阵，不做详述
  • PointNet中有k*k的T-net，和3*3的T-net没有太大区别，主要对于输入数据的处理上一点点不同，不做详述

PointNet++