PointNet笔记

可能遇到的问题

在windows上运行pointnet的代码时，可能会遇到一些问题：
1.比如提示OSError: no file with expected extension， 这是因为可视化的show3d_balls.py文件运行不了， 具体的解决方法可以看这篇文章：https://blog.51cto.com/u_16213693/7738038。
2.由于作者的pointnet所用的pytorch版本比较低，可能会出现：AttributeError: module ‘distutils‘ has no attribute ‘version的报错，具体的解决方法可以参考https://blog.csdn.net/qq_42076902/article/details/129261266。

代码详解

点云数据集加载部分不在这里过多解释，具体的注释代码请见

    从网络结构图可以知道，分割网络其实也用到了分类网络的一部分。拿分割部分来说，程序的入口在文件 train_segmentation.py 中，其中seger = PointNetDenseCls(k=num_classes, feature_transform=opt.feature_transform) 这段代码代表了初始化一个分割模型，接下来我们进入到这个模型中去查看具体细节。
    按住Ctrl＋指针点击PointNetDenseCls，进入到 model.py 中，再看PointNetDenseCls这个类之前，我们先根据这个文件中定义的类把整个网络架构梳理一遍：首先我们先看图中的蓝色部分，也就是分类网络，这时进入到PointNetfeat这个类中查看，代码和注释如下：

# 这是一个三维旋转变换网络
class STN3d(nn.Module):
    def __init__(self):
        super(STN3d, self).__init__()
        # 使用大小为1*1的卷积核进行升维
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 9)
        self.relu = nn.ReLU()

        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)

    # 网络向前传播
    def forward(self, x):
        # 获取点的数量，假设输入的x是(32, 3, 2500)
        batchsize = x.size()[0] # 32
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x))) # shape:(32, 1024, 2500)
        x = torch.max(x, 2, keepdim=True)[0] # shape:(32, 1024, 1)
        x = x.view(-1, 1024) # (32, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)
        # 下面这段代码用于计算旋转平移矩阵，但是后续的论文证明这个方法并无太大意义，故不做太多解释
        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
        x = x.view(-1, 3, 3)
        return x

# k维度旋转变换网络，与上面的大同小异。
class STNkd(nn.Module):
    def __init__(self, k=64):
        super(STNkd, self).__init__()
        self.conv1 = torch.nn.Conv1d(k, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k * k)
        self.relu = nn.ReLU()

        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)

        self.k = k

    def forward(self, x):
        batchsize = x.size()[0]
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)

        iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1, self.k * self.k).repeat(
            batchsize, 1)
        if x.is_cuda:
            iden = iden.cuda()
        x = x + iden
        x = x.view(-1, self.k, self.k)
        return x


# 分类网络的前面的部分
class PointNetfeat(nn.Module):
    def __init__(self, global_feat=True, feature_transform=False):
        # 表示调用PointNetfeat的父类nn.Module的init方法
        super(PointNetfeat, self).__init__()
        # 定义一个变换网络
        self.stn = STN3d()
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.global_feat = global_feat
        self.feature_transform = feature_transform
        # 是否要用到特征变换网络？
        if self.feature_transform:
            self.fstn = STNkd(k=64)

    def forward(self, x): #(32, 3, 2500)
        n_pts = x.size()[2] # 2500
        trans = self.stn(x) # 学习一个特征变换网络
        x = x.transpose(2, 1) # (32, 2500, 3)
        x = torch.bmm(x, trans) # 批量矩阵乘法，trans为(32, 3, 3)
        x = x.transpose(2, 1) # 再交换回来 (32, 3, 2500)
        x = F.relu(self.bn1(self.conv1(x)))

        # 是否运用到特征变换？ 如果用了，下面的步骤与前面的逻辑相同
        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)))
        x = self.bn3(self.conv3(x))
        # 应用最大池化操作，变为(32, 1024, 1)
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)
        if self.global_feat:
            return x, trans, trans_feat
        else:
            x = x.view(-1, 1024, 1).repeat(1, 1, n_pts) # 变为 (32, 1024, 2500)
            return torch.cat([x, pointfeat], 1), trans, trans_feat


# 分类网络的最终部分
class PointNetCls(nn.Module):
    def __init__(self, k=2, feature_transform=False):
        super(PointNetCls, self).__init__()
        self.feature_transform = feature_transform
        self.feat = PointNetfeat(global_feat=True, feature_transform=feature_transform)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k)
        self.dropout = nn.Dropout(p=0.3)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.relu = nn.ReLU()

    def forward(self, x):
        x, trans, trans_feat = self.feat(x)
        x = F.relu(self.bn1(self.fc1(x)))
        x = F.relu(self.bn2(self.dropout(self.fc2(x))))
        x = self.fc3(x)
        return F.log_softmax(x, dim=1), trans, trans_feat


# 这个类是分割部分总架构
class PointNetDenseCls(nn.Module):
    def __init__(self, k=2, feature_transform=False):
        super(PointNetDenseCls, self).__init__()
        self.k = k # k代表类别个数，默认值是2
        self.feature_transform = feature_transform
        self.feat = PointNetfeat(global_feat=False, feature_transform=feature_transform)
        self.conv1 = torch.nn.Conv1d(1088, 512, 1) # 1024 + 64
        self.conv2 = torch.nn.Conv1d(512, 256, 1)
        self.conv3 = torch.nn.Conv1d(256, 128, 1)
        self.conv4 = torch.nn.Conv1d(128, self.k, 1)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.bn3 = nn.BatchNorm1d(128)

    def forward(self, x):
        batchsize = x.size()[0]
        n_pts = x.size()[2]
        x, trans, trans_feat = self.feat(x)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.conv4(x)
        x = x.transpose(2, 1).contiguous()
        x = F.log_softmax(x.view(-1, self.k), dim=-1) # 表示在最后一个维度上应用 对数softmax函数
        x = x.view(batchsize, n_pts, self.k)
        return x, trans, trans_feat

train_segmentation.py的注释

from __future__ import print_function
import argparse
import sys

sys.path.append("../")
import os
import random
import torch
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
from pointnet.dataset import ShapeNetDataset
from pointnet.model import PointNetDenseCls, feature_transform_regularizer
import torch.nn.functional as F
from tqdm import tqdm
import numpy as np


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


if __name__ == '__main__':
    # 创建命令行解析器
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--batchSize', type=int, default=16, help='input batch size')
    parser.add_argument(
        '--workers', type=int, help='number of data loading workers', default=4)
    parser.add_argument(
        '--nepoch', type=int, default=25, help='number of epochs to train for')
    parser.add_argument('--outf', type=str, default='seg', help='output folder')
    parser.add_argument('--model', type=str, default='', help='model path')
    parser.add_argument('--dataset', type=str, required=True, help="dataset path")
    parser.add_argument('--class_choice', type=str, default='Chair', help="class_choice")
    parser.add_argument('--feature_transform', action='store_true', help="use feature transform")

    # 解析输入的命令行参数，并存储在opt中
    opt = parser.parse_args()
    print(opt)

    # 为opt对象添加一个manualSeed属性，将其设置为1-10000之间的一个随机值。
    opt.manualSeed = random.randint(1, 10000)
    # 这一行代码使用生成的随机种子来初始化 Python 内置的 random 模块的随机数生成器
    random.seed(opt.manualSeed)
    # 这一行代码使用生成的随机种子来初始化 PyTorch 的随机数生成器
    torch.manual_seed(opt.manualSeed)

    # 创建一个数据集实例(进入到ShapeNetDataset)中继续看
    dataset = ShapeNetDataset(
        root=opt.dataset,
        classification=False,
        class_choice=[opt.class_choice])
    # 数据集加载器
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=opt.batchSize,
        shuffle=True,
        num_workers=int(opt.workers))

    test_dataset = ShapeNetDataset(
        root=opt.dataset,
        classification=False,
        class_choice=[opt.class_choice],
        split='test',
        data_augmentation=False)
    testdataloader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=opt.batchSize,
        shuffle=True,
        num_workers=int(opt.workers))

    print(len(dataset), len(test_dataset))
    num_classes = dataset.num_seg_classes
    print('classes', num_classes)
    try:
        os.makedirs(opt.outf)
    except OSError:
        pass

    blue = lambda x: '\033[94m' + x + '\033[0m'

    # 初始化分割模型，PointNetDenseCls为模型
    seger = PointNetDenseCls(k=num_classes, feature_transform=opt.feature_transform)

    # 是否有预训练的模型？
    if opt.model != '':
        seger.load_state_dict(torch.load(opt.model))

    # 使用Adam优化算法创建一个优化器
    optimizer = optim.Adam(seger.parameters(), lr=0.001, betas=(0.9, 0.999))

    # 创建学习率调度器
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
    seger.cuda()

    num_batch = len(dataset) / opt.batchSize

    # 开始训练模型
    for epoch in range(opt.nepoch):
        scheduler.step()
        for i, data in enumerate(dataloader, 0):
            # 从数据中解析出点和标签
            points, target = data
            # 交换维度，便于后续操作
            points = points.transpose(2, 1)
            # 利用GPU来学习
            points, target = points.cuda(), target.cuda()
            # 清零优化器中所有模型参数的梯度。这是因为梯度在 PyTorch 中默认是累加的。
            optimizer.zero_grad()
            # 将模型设置为训练模式，这将启用模型中的训练相关操作，如 Dropout。
            seger = seger.train()
            # 进行向前传播，得到预测值等信息。
            pred, trans, trans_feat = seger(points)
            pred = pred.view(-1, num_classes)
            target = target.view(-1, 1)[:, 0] - 1
            # 计算损失，使用负对数似然损失
            loss = F.nll_loss(pred, target)
            if opt.feature_transform:
                loss += feature_transform_regularizer(trans_feat) * 0.001
            # 反向传播
            loss.backward()
            # 使用优化器更新模型参数，根据计算的梯度进行优化。
            optimizer.step()
            # 得到每个样本预测类别的索引
            pred_choice = pred.data.max(1)[1]
            # 得到正确预测的样本数量
            correct = pred_choice.eq(target.data).cpu().sum()
            print('[%d: %d/%d] train loss: %f accuracy: %f' % (
                epoch, i, num_batch, loss.item(), correct.item() / float(opt.batchSize * 2500)))

            # 测试部分
            if i % 10 == 0:
                j, data = next(enumerate(testdataloader, 0))
                points, target = data
                points = points.transpose(2, 1)
                points, target = points.cuda(), target.cuda()
                seger = seger.eval()
                pred, _, _ = seger(points)
                pred = pred.view(-1, num_classes)
                target = target.view(-1, 1)[:, 0] - 1
                loss = F.nll_loss(pred, target)
                pred_choice = pred.data.max(1)[1]
                correct = pred_choice.eq(target.data).cpu().sum()
                print('[%d: %d/%d] %s loss: %f accuracy: %f' % (
                    epoch, i, num_batch, blue('test'), loss.item(), correct.item() / float(opt.batchSize * 2500)))

        torch.save(seger.state_dict(), '%s/seg_model_%s_%d.pth' % (opt.outf, opt.class_choice, epoch))

# benchmark mIOU
    # 初始化一个空列表，用于存储每个形状的Iou值
    shape_ious = []
    for i, data in tqdm(enumerate(testdataloader, 0)):
        points, target = data
        points = points.transpose(2, 1)
        points, target = points.cuda(), target.cuda()
        # 将模型seger设置为评估模式，以确保模型在推理时的行为与训练时不同，例如关闭dropout等
        seger = seger.eval()
        pred, _, _ = seger(points)
        pred_choice = pred.data.max(2)[1]

        pred_np = pred_choice.cpu().data.numpy()
        target_np = target.cpu().data.numpy() - 1

        for shape_idx in range(target_np.shape[0]):
            parts = range(num_classes)  # np.unique(target_np[shape_idx])
            part_ious = []
            for part in parts:
                I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part))
                U = np.sum(np.logical_or(pred_np[shape_idx] == part, target_np[shape_idx] == part))
                if U == 0:
                    iou = 1  # If the union of groundtruth and prediction points is empty, then count part IoU as 1
                else:
                    iou = I / float(U)
                part_ious.append(iou)
            shape_ious.append(np.mean(part_ious))

    print("mIOU for class {}: {}".format(opt.class_choice, np.mean(shape_ious)))

dataset.py文件中的shapenetdataset类的注释

# ShapeNetDataset继承自data.Dataset
class ShapeNetDataset(data.Dataset):
    def __init__(self,
                 root,
                 npoints=2500,
                 classification=False,
                 class_choice=None,
                 split='train',
                 data_augmentation=True):
        self.npoints = npoints
        self.root = root
        self.catfile = os.path.join(self.root, 'synsetoffset2category.txt')
        self.cat = {}
        self.data_augmentation = data_augmentation
        self.classification = classification
        self.seg_classes = {}

        # 注意这里的class_choice默认是Chair

        with open(self.catfile, 'r') as f:
            for line in f:
                # 按照空格分割  例如:Chair 03001627
                ls = line.strip().split()
                # 字典 key value
                self.cat[ls[0]] = ls[1]

        # 这是一个字典推导式，过滤 self.cat 字典，只保留那些键（k）在 class_choice 列表中的键值对。在这里，只取Chair
        if class_choice is not None:
            self.cat = {k: v for k, v in self.cat.items() if k in class_choice}

        # 键值对互换  03001627 Chair
        self.id2cat = {v: k for k, v in self.cat.items()}

        self.meta = {}
        splitfile = os.path.join(self.root, 'train_test_split', 'shuffled_{}_file_list.json'.format(split))
        # 从json文件中解析数据
        filelist = json.load(open(splitfile, 'r'))
        # 默认情况下item为键
        for item in self.cat:
            self.meta[item] = []

        for file in filelist:
            _, category, uuid = file.split('/') #例如 _ ,  03001627, 7fe08cd7a9b76c1dcbde89e0c48a01bf
            if category in self.cat.values():
                self.meta[self.id2cat[category]].append((os.path.join(self.root, category, 'points', uuid + '.pts'),
                                                         os.path.join(self.root, category, 'points_label',
                                                                      uuid + '.seg')))

        self.datapath = []
        # 下面的item是名称而不是数字
        for item in self.cat:
            for fn in self.meta[item]:
                self.datapath.append((item, fn[0], fn[1]))
        # zip(['Chair', 'Lamp', 'Table'], [0, 1, 2])
        # ->
        # [('Chair', 0), ('Lamp', 1), ('Table', 2)]
        # ->
        # {
        #     'Chair': 0,
        #     'Lamp': 1,
        #     'Table': 2
        # }
        # 以上是下面这行代码的注释, 最终classes是一个字典，结果为{'Chair' : 0}
        self.classes = dict(zip(sorted(self.cat), range(len(self.cat))))
        # print(self.classes)

        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/num_seg_classes.txt'), 'r') as f:
            for line in f:
                ls = line.strip().split()
                self.seg_classes[ls[0]] = int(ls[1])
        self.num_seg_classes = self.seg_classes[list(self.cat.keys())[0]]
        print(self.seg_classes, self.num_seg_classes)

    def __getitem__(self, index):
        fn = self.datapath[index]
        cls = self.classes[self.datapath[index][0]]
        point_set = np.loadtxt(fn[1]).astype(np.float32)
        seg = np.loadtxt(fn[2]).astype(np.int64)
        # print(point_set.shape, seg.shape)

        # 重新采样点集，数量为self.npoints
        choice = np.random.choice(len(seg), self.npoints, replace=True)
        # 注意这里point_set是二维数组，第二维为3
        point_set = point_set[choice, :]

        # 将点云进行中心化处理，中心化是将点云质心移动到原点， 注意这里axis=0意为列，并不是行，np.mean(point_set, axis=0)
        # 计算的是每一列的均值，结果是一个形状为 (D,) 的数组(原数组为(N,D))，np.expand_dims(..., 0) 将这个均值数组扩展一个维度，
        # 变成一个形状为 (1, D) 的数组。这个操作是为了与 point_set 进行广播运算。结果是每个点都减去了整体的质心，使得新的点云数据中心在原点。
        point_set = point_set - np.expand_dims(np.mean(point_set, axis=0), 0)

        # np.sum(point_set ** 2, axis=1) 计算每个点的平方和，结果是一个形状为 (N,) 的数组，其中 N 是点的数量。例如，如果 point_set
        # 中的一个点是 [x, y, z]，则平方和为 x^2 + y^2 + z^2。
        # np.sqrt(...) 计算这些平方和的平方根，得到每个点到原点的欧氏距离。
        # np.max(...) 找到这些距离中的最大值, 原作者在np.max函数中加了个0，但是这应该是不必要的，去掉。
        dist = np.max(np.sqrt(np.sum(point_set ** 2, axis=1)))

        # 归一化点云数据 point_set / dist 将 point_set 中的每个点除以最大距离 dist，结果是将所有点缩放到单位球内，使得最大距离为1。
        point_set = point_set / dist

        # 数据增强，通过随机旋转和抖动点云数据，使模型在训练过程中能够见到更多样化的数据，从而提高模型的泛化能力。
        if self.data_augmentation:
            # 生成随机旋转角度，范围为(0, 2π)
            theta = np.random.uniform(0, np.pi * 2)
            # 通过生成的随机角度 theta 构建一个二维旋转矩阵。旋转矩阵用于在二维平面上旋转点。这里的旋转是针对 x 和 z 轴进行的
            # （在三维空间中忽略了 y 轴），因为点云的旋转一般在水平面上更常见。
            rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
            # 选择点云数据的 x 和 z 维度（即列索引 0 和 2），并对它们应用旋转矩阵进行旋转。
            # point_set[:, [0, 2]] 选择 point_set 中所有点的 x 和 z 坐标。
            # dot(rotation_matrix) 对选中的坐标应用旋转矩阵进行旋转。
            point_set[:, [0, 2]] = point_set[:, [0, 2]].dot(rotation_matrix)  # random rotation
            # 通过向点云数据添加随机噪声，进一步增强数据的多样性。
            # np.random.normal(0, 0.02, size=point_set.shape) 生成与 point_set 形状相同的随机噪声，噪声服从均值为 0、标准差为 0.02 的正态分布。
            # point_set += ... 将生成的随机噪声添加到原始点云数据中。
            point_set += np.random.normal(0, 0.02, size=point_set.shape)  # random jitter

        seg = seg[choice]
        point_set = torch.from_numpy(point_set)
        seg = torch.from_numpy(seg)
        cls = torch.from_numpy(np.array([cls]).astype(np.int64))

        if self.classification:
            return point_set, cls
        else:
            return point_set, seg

    def __len__(self):
        return len(self.datapath)