PointNet代码解读
引用说明
文章中的代码全部来自于Github仓库: Pointnet_Pointnet2_pytorch
本文更关注语义分割以及零件分割部分代码
PointNet
- 完整的PointNet结构代码
来自于models文件夹下的pointnet_utils.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_utils.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
- PointNet零件分割任务
来自models下面的pointnet_cls.py文件,这里并没有复用之前的模型,而是自己实现了一个,某几个地方有区别都做了注释
首先是get_model函数
class get_model(nn.Module):
def __init__(self, part_num=50, normal_channel=True): # 这里由于用的partnet的数据集,一共有50个部件标号所以初始化为50
super(get_model, self).__init__()
if normal_channel: # 是否有法向量
channel = 6
else:
channel = 3
self.part_num = part_num
self.stn = STN3d(channel) # 输入数据的旋转矩阵
# 五个一维卷积用于升维度
self.conv1 = torch.nn.Conv1d(channel, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 128, 1)
self.conv4 = torch.nn.Conv1d(128, 512, 1)
self.conv5 = torch.nn.Conv1d(512, 2048, 1)
# bn
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(128)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(2048)
self.fstn = STNkd(k=128) # 这里用的是128*128的旋转矩阵
self.convs1 = torch.nn.Conv1d(4944, 256, 1) # 4944是怎么来的?
self.convs2 = torch.nn.Conv1d(256, 256, 1)
self.convs3 = torch.nn.Conv1d(256, 128, 1)
self.convs4 = torch.nn.Conv1d(128, part_num, 1) # 输出到50个类别
self.bns1 = nn.BatchNorm1d(256)
self.bns2 = nn.BatchNorm1d(256)
self.bns3 = nn.BatchNorm1d(128)
def forward(self, point_cloud, label):
B, D, N = point_cloud.size()
trans = self.stn(point_cloud)
point_cloud = point_cloud.transpose(2, 1)
if D > 3:
point_cloud, feature = point_cloud.split(3, dim=2)
point_cloud = torch.bmm(point_cloud, trans)
if D > 3:
point_cloud = torch.cat([point_cloud, feature], dim=2)
point_cloud = point_cloud.transpose(2, 1)
out1 = F.relu(self.bn1(self.conv1(point_cloud))) # 第一个MLP 3->64
out2 = F.relu(self.bn2(self.conv2(out1))) # 64 -> 128
out3 = F.relu(self.bn3(self.conv3(out2))) # 128 -> 128 为什么做一个这个?
trans_feat = self.fstn(out3) # 特征矩阵的旋转矩阵 128*128
x = out3.transpose(2, 1)
net_transformed = torch.bmm(x, trans_feat)
net_transformed = net_transformed.transpose(2, 1)
out4 = F.relu(self.bn4(self.conv4(net_transformed))) # 这里 238 -> 512
out5 = self.bn5(self.conv5(out4)) # 512 -> 2048
out_max = torch.max(out5, 2, keepdim=True)[0] # 这里max-pooling
out_max = out_max.view(-1, 2048) # 整理维度
out_max = torch.cat([out_max,label.squeeze(1)],1) # squezze() 删除维度为1的维度
expand = out_max.view(-1, 2048+16, 1).repeat(1, 1, N) # +16这个16是一共有16种数据
concat = torch.cat([expand, out1, out2, out3, out4, out5], 1) # 由于做零件分割,他这里不是只连接了一层的feature,而是将每一次MLP的都连接进来了 4944=64+128+128+512+2048+(2048+16)
net = F.relu(self.bns1(self.convs1(concat)))
net = F.relu(self.bns2(self.convs2(net)))
net = F.relu(self.bns3(self.convs3(net)))
net = self.convs4(net) # 到这里已经输出到50个类别了
net = net.transpose(2, 1).contiguous()
net = F.log_softmax(net.view(-1, self.part_num), dim=-1) # 这里的log_softmax就是softmax之后取log 因此有正有负
net = net.view(B, N, self.part_num) # [B, N, 50]
return net, trans_feat
然后是get_loss函数
class get_loss(torch.nn.Module):
def __init__(self, mat_diff_loss_scale=0.001): # 这个scale是为了调整正则项的权重
super(get_loss, self).__init__()
self.mat_diff_loss_scale = mat_diff_loss_scale
def forward(self, pred, target, trans_feat):
loss = F.nll_loss(pred, target) # nll_loss下文单独讲下
mat_diff_loss = feature_transform_reguliarzer(trans_feat) # 这个就是文章中说的那个正则项
total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale # 加权重与nll_loss求和拼接成损失函数
return total_loss
- 补充
feature_transform_reguliarzer
方法主要是对于T-net生成的矩阵正则的处理,希望是更接近正交矩阵,不做详述- PointNet中有k*k的T-net,和3*3的T-net没有太大区别,主要对于输入数据的处理上一点点不同,不做详述
- log_softmax & nll_loss
- 首先log_softmax就是在softmax基础上取log,但是实际实现方法不是这样
- 举例说明下,输入分别为预测的每个类别可能性的向量以及他们真实的标号(label)
input = tensor([[-0.3301, 0.0966, 1.5706, -0.3922, 0.3910], [ 1.6564, -0.1009, 1.0399, 0.4459, -0.2528], [ 0.6746, 1.8496, -0.0700, 0.3501, 0.4426], [-2.0736, -0.3385, 0.9388, 1.6099, -0.4514]]) # input N*C N是数据量 C是分类个数 这里是4*5 target = torch.tensor([4, 1, 3, 0]) # 最大值小于C 最小值大于等于0
- 验证
output = F.nll_loss(F.log_softmax(input, dim=1), target) # 分别输入log_softmax结果和target # tensor(1.9301) output1 = torch.log(F.softmax(input)) tensor([[-2.7174, -1.8518, -0.4467, -3.5710, -2.2152], [-0.8904, -0.9595, -2.4950, -3.5999, -2.3370], [-2.7413, -1.3464, -4.6853, -0.6783, -1.8413], [-3.8673, -2.4171, -0.5103, -2.9536, -1.4378]]) # nll_loss就是按照target的序号取出output1中对应的数据 -2.2152-0.9595-0.6783-3.8673 = -7.7203 / 4 = -1.9301 去除负号, reduction参数用来定义是对于-7.7203取平均还是求和 mean or sum