基于PointNet网络

编辑中

表达旋转：四元数

原理简析

详细可以看SLAM14讲关于四元数的部分，如果看不懂，可以看相关资料中的虚数视频。 四元数在点云中是一个很好的表示旋转的方式。 一个三维空间点用四元数表示如下

$$P = [0, x, y, z] = [0, v]$$

三维空间的点是个纯虚四元数（实部为0），三个虚部与空间中的三个轴对应。

怎么用四元数表示一个旋转呢？如下所示

$$q = \left [ \cos \frac{\theta }{2}， n\sin \frac{\theta }{2} \right ]$$

左边的cos是实部，也就是说，实部存储的是旋转角度 右边的sin是虚部，n表示的是旋转轴，旋转轴有三个维度，所以虚部有三个数，每个数都要分别乘以一个θ/2。也就是说，虚部存放的是旋转轴，表示物体绕着哪个轴进行旋转

那么，一个点绕着一个轴进行旋转，旋转后的点怎么计算出来呢？

$$p' = qpq^{-1}$$

其中，p'是结果，通过原点p左乘旋转四元数，右乘旋转四元数的转置得到，结果就是新的点p的四元数数值

一个单位四元数（即模长为1的四元数）可以表示三维空间中的任意旋转，

用Numpy实现四元数运算

import numpy as np

x = np.radians(90)  # 绕x轴进行90度运算
x_vector = np.array([1, 0, 0])  # 定义旋转轴，x轴的单位向量

def quaternion_multiply(quaternion1, quaternion0):  # 四元数相乘函数
    w0, x0, y0, z0 = quaternion0
    w1, x1, y1, z1 = quaternion1
    return np.array([-x1 * x0 - y1 * y0 - z1 * z0 + w1 * w0,
                     x1 * w0 + y1 * z0 - z1 * y0 + w1 * x0,
                     -x1 * z0 + y1 * w0 + z1 * x0 + w1 * y0,
                     x1 * y0 - y1 * x0 + z1 * w0 + w1 * z0], dtype=np.float64)

def quaternion_contrary(quaternion):  # 四元数取逆函数
    w, x, y, z = quaternion
    q_star = np.array([w, -x, -y, -z])  # 共轭
    q_abs_pow2 = np.sum(np.power(quaternion, 2)) # 模长的平方
    return q_star / q_abs_pow2

p = np.array([0, 0, 2, 0])  # 定义一个点，在y轴上 (0, 2, 0) 如果此点绕x轴旋转，会变成 (0, 0, 2)
q_x = np.hstack((np.cos([x / 2]), np.sin(x / 2) * x_vector))  # 旋转四元数 (cosx/2, nsinx/2)
print(q_x)
print(quaternion_multiply(quaternion_multiply(q_x, p), quaternion_contrary(q_x)))python

最终输出的结果为

[0.70710678 0.70710678 0.         0.        ]
[0.0000000e+00 0.0000000e+00 4.4408921e-16 2.0000000e+00]

第一行，这是一个四元数，表示绕x轴旋转90度。这个四元数的实部为cos(旋转角/2)，虚部的第一个元素为1*sin(旋转角/2)，其余元素为0，符合预期，虚部的三个元素表示旋转轴(1, 0, 0)的方向乘以sin(旋转角度/2)。 第二行，表示结果的四元数表示，第一维度为0，也就是实部为0，表示是一个点。三个虚部表示点的坐标，其中第二个虚部非常小，可以看作是0，所以得到的坐标为(0, 0, 2)，符合预期。

解决思路

学习旋转角度，也就是分别绕着x轴、y轴、z轴旋转的角度，跟欧拉角差不多，但是用四元数进行计算，能够避免欧拉角的奇异性问题。

损失函数：两个四元数a和b，a乘以b的逆，能得到一个新的四元数，新的四元数表示了一个绕着新的轴进行的一次a到b的相对旋转，也就是说损失是拿这个新的四元数的实部的旋转角，旋转角越小表示两个旋转四元数越接近

网络结构

源码

四元数计算类 Quaternion

import torch
class Quaternion:
    def __init__(self, w, x, y, z):
        self.value = torch.tensor([w, x, y, z])

    def multiply(self, quaternion):  # 与另一个四元数相乘
        w0, x0, y0, z0 = self.value
        w1, x1, y1, z1 = quaternion
        return torch.tensor([-x1 * x0 - y1 * y0 - z1 * z0 + w1 * w0,
                         x1 * w0 + y1 * z0 - z1 * y0 + w1 * x0,
                         -x1 * z0 + y1 * w0 + z1 * x0 + w1 * y0,
                         x1 * y0 - y1 * x0 + z1 * w0 + w1 * z0])

    def conjugate(self):  # 取共轭数
        w, x, y, z = self.value
        q_star = torch.tensor([w, -x, -y, -z])  # 共轭
        return q_star

    def contrary(self):  # 取转置
        q_star = self.conjugate()
        q_abs_pow2 = torch.sum(self.value**2)  # 模长的平方
        return q_star / q_abs_pow2

    def getRotationMatrix(self):  # 从四元数转旋转矩阵
        q0, q1, q2, q3 = self.value
        return torch.tensor([[1 - 2*q2*q2 - 2*q3*q3, 2*q1*q2 + 2*q0*q3, 2*q1*q3 - 2*q0*q2],
                         [2*q1*q2 - 2*q0*q3, 1-2*q1*q1 - 2*q3*q3, 2*q2*q3 + 2*q0*q1],
                         [2*q1*q3 + 2*q0*q2, 2*q2*q3 - 2*q0*q1, 1-2*q1*q1-2*q2*q2]])

    def getValue(self):
        return self.value

    @staticmethod
    def from_tensor(array):  # 从tensor转成四元数类
        array = array.squeeze()
        return Quaternion(array[0], array[1], array[2], array[3])

    @staticmethod
    def get_rotation_quaternion(axis, angle):  # 获得旋转四元数。0:x轴 1:y轴 2:z轴, 旋转角度
        angle = torch.tensor(angle, dtype=torch.float32) * (torch.pi / 180)

        R_vector = None
        if axis == 0:
            R_vector = torch.tensor([1, 0, 0])
        elif axis == 1:
            R_vector = torch.tensor([0, 1, 0])
        elif axis == 2:
            R_vector = torch.tensor([0, 0, 1])

        return Quaternion.from_tensor(torch.cat((torch.cos(angle / 2), torch.sin(angle / 2) * R_vector)))

数据预处理 Dataset

import random
from torch.utils.data import dataset
import open3d as o3d
import os
import numpy as np
import Quaternion as quat
import torch


def random_unit_quaternion(transform_range):
    # 随机生成旋转角（0 到 180 度），并将其转换为弧度
    angle = torch.rand(1) * torch.pi * 2
    # 计算实部
    w = torch.cos(angle / 2)

    # 随机生成虚部
    x, y, z = 2 * torch.rand(3) - 1
    v_norm = torch.sqrt(x ** 2 + y ** 2 + z ** 2)

    # 单位化虚部向量
    x, y, z = x / v_norm, y / v_norm, z / v_norm

    # 计算旋转轴的缩放系数
    s = torch.sin(angle / 2)

    # 根据实部和虚部组合成单位四元数
    q = torch.tensor([w, x * s, y * s, z * s])
    if w < 0:
        q = -q  # 始终保持第一个数是正数

    # 将四元数转换为旋转矩阵
    quat_obj = quat.Quaternion.from_tensor(q)
    R = quat_obj.getRotationMatrix()
    t = 2 * transform_range * torch.rand(3) - transform_range
    return quat_obj, R, t


class LoadData(dataset.Dataset):
    def __init__(self, points_count, std_model_paths, t_range,label_is_rotation):
        super(LoadData, self).__init__()
        folder_path = std_model_paths
        self.file_paths = []
        self.points_count = points_count
        self.count = 0
        self.pcds_array = []
        self.label_is_rotation = label_is_rotation
        self.t_range = t_range

        for root, dirs, files in os.walk(folder_path):
            for file in files:
                if file == '.DS_Store':
                    continue
                self.count = self.count + 1
                self.pcds_array.append(torch.tensor(np.asarray(o3d.io.read_point_cloud(os.path.join(root, file)).points), dtype=torch.float32))
                print(os.path.join(root, file))

    def __getitem__(self, index):
        with torch.no_grad():
            pcd_std = self.pcds_array[index % self.count]  # 标准点云
            random_index = random.randint(0, self.count-1)  # 生成一个0到self.count之间的随机整数
            pcd_aligned = self.pcds_array[random_index]  # 被对其点云
            q, R, t = random_unit_quaternion(self.t_range)  # 获取随机旋转四元数
            sampled_indices_std = torch.randperm(pcd_std.shape[0])[:self.points_count]
            pcd_std = pcd_std[sampled_indices_std]  # 从模型中随机采样n个点
            sampled_indices_aligned = torch.randperm(pcd_aligned.shape[0])[:self.points_count]
            pcd_aligned = pcd_aligned[sampled_indices_aligned]  # 从模型中随机采样n个点
            points_after = torch.matmul(pcd_aligned - t, R.T)  # 将采样到的点运用 随机旋转四元数
            if self.label_is_rotation:
                label = q.getValue()
            else:
                label = t
        return points_after, pcd_std, label # 返回随机旋转后的点云 以及对应的四元数

    def __len__(self):
        return 500  # self.count

1. 随机生成单位四元数（random_unit_quaternion）： 这个函数随机生成一个单位四元数表示的旋转，并将其转换为旋转矩阵。同时，生成一个随机平移向量。这个函数的主要作用是为点云数据应用随机的旋转和平移。

2. 点云数据集类（LoadData）： 这个类继承自torch.utils.data.Dataset，用于加载点云数据并进行预处理。主要功能如下：

   • 在__init__方法中，读取指定文件夹下的所有点云文件，并将其转换为Tensor格式存储。

   • 在__getitem__方法中，对点云数据进行以下处理： a. 随机选择两个点云（标准点云和被对齐点云）。 b. 随机生成旋转四元数和平移向量。 c. 对点云数据进行随机采样，选取指定数量的点。 d. 将被对齐点云应用随机旋转和平移，得到旋转后的点云。 e. 返回旋转后的点云、标准点云和对应的四元数或平移向量（根据label_is_rotation参数决定）。

   • 在__len__方法中，返回数据集的长度，这里设置为固定值500。

定义网络 Net

import torch
import torch.nn as nn
import numpy as np


class PointNet(nn.Module):
    def __init__(self):
        super(PointNet, self).__init__()
        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)


    def forward(self, x):
        x = torch.relu(self.bn1(self.conv1(x)))  # 1D卷积层 3 -> 64、归一化、Relu
        x = torch.relu(self.bn2(self.conv2(x)))  # 1D卷积层 64->128、归一化、Relu
        x = self.bn3(self.conv3(x))  # 1D卷积层 128->1024、归一化
        x = torch.max(x, 2, keepdim=True)[0]  # 最大池化 （batch_size, 1024）
        x = x.view(-1, 1024)
        return x


class RotationNet(nn.Module):
    def __init__(self):
        super(RotationNet, self).__init__()
        self.pointnet = PointNet()
        self.fc1 = nn.Linear(2048, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 4)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(64)

    def forward(self, x1, x2):
        # 计算x1和x2的质心
        centroid_x1 = torch.mean(x1, dim=2, keepdim=True)  # (32, 3, 1)
        centroid_x2 = torch.mean(x2, dim=2, keepdim=True)  # (32, 3, 1)

        # 将x1和x2归一化到质心位置
        x1 = x1 - centroid_x1
        x2 = x2 - centroid_x2

        # 特征提取
        x1 = self.pointnet(x1)
        x2 = self.pointnet(x2)

        x = torch.cat((x1, x2), dim=1)
        x = torch.relu(self.bn1(self.fc1(x)))
        x = torch.relu(self.bn2(self.fc2(x)))
        x = self.fc3(x).float()
        x = torch.nn.functional.normalize(x, p=2, dim=1)
        # 确保四元数第一个分量为正值
        mask = x[:, 0] < 0
        x[mask] = -x[mask]
        return x


class TranslationNet(nn.Module):
    def __init__(self):
        super(TranslationNet, self).__init__()
        self.pointnet = PointNet()
        self.fc1 = nn.Linear(2048, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 3)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(64)

    def forward(self, x1, x2):
        x1 = self.pointnet(x1)
        x2 = self.pointnet(x2)
        x = torch.cat((x1, x2), dim=1)
        x = torch.relu(self.bn1(self.fc1(x)))
        x = torch.relu(self.bn2(self.fc2(x)))
        x = self.fc3(x).float()
        return x

1. PointNet（点云特征提取网络）:

   • 包含三个1D卷积层和相应的批归一化层。

   • 在forward方法中，输入的点云数据依次经过三个1D卷积层和批归一化层，然后通过ReLU激活函数。

   • 最后，对结果进行最大池化操作，得到大小为(batch_size, 1024)的特征表示。

2. RotationNet（旋转预测网络）：

   • 包含一个PointNet实例用于特征提取，以及三个全连接层和相应的批归一化层。

   • 在forward方法中，首先计算输入点云的质心，然后将点云归一化到质心位置。

   • 使用PointNet提取点云特征，然后将两个点云的特征沿维度1拼接。

   • 拼接后的特征依次通过全连接层、批归一化层和ReLU激活函数。

   • 最后，经过一层全连接层并进行L2范数归一化，得到大小为(batch_size, 4)的四元数表示。

   • 保证四元数的第一个分量为正值。

3. TranslationNet（平移预测网络）：

   • 包含一个PointNet实例用于特征提取，以及三个全连接层和相应的批归一化层。

   • 在forward方法中，使用PointNet提取点云特征，然后将两个点云的特征沿维度1拼接。

   • 拼接后的特征依次通过全连接层、批归一化层和ReLU激活函数。

   • 最后，经过一层全连接层，得到大小为(batch_size, 3)的平移向量表示。

四元数损失定义 QuatLoss

import torch
import torch.nn as nn
import numpy as np


class QuatLoss(nn.Module):
    def __init__(self):
        super(QuatLoss, self).__init__()

    def forward(self, y_hat, label):
        w, x, y, z = label.split(1, dim=-1)
        label_inv = torch.cat([w, -x, -y, -z], dim=-1)  # label共轭
        sa, xa, ya, za = y_hat.split(1, dim=-1)
        sb, xb, yb, zb = label_inv.split(1, dim=-1)
        loss_quat = torch.cat([
            sa * sb - xa * xb - ya * yb - za * zb,
            sa * xb + xa * sb + ya * zb - za * yb,
            sa * yb - xa * zb + ya * sb + za * xb,
            sa * zb + xa * yb - ya * xb + za * sb
        ], dim=-1)
        # loss = torch.pow(torch.acos(loss_quat[:, 0]) * 2, 2)
        cos_angle = loss_quat[:, 0]  # 四元数点积（它们之间的夹角余弦值）
        loss = 1 - torch.square(cos_angle)  # 计算 1 减去余弦值的平方

        return torch.mean(loss)

这段代码定义了一个四元数损失函数（QuatLoss），它继承自torch.nn.Module。这个损失函数的目的是计算两个四元数之间的角度差异，并使用这个差异作为预测旋转的损失。

在forward方法中：

1. 将标签（label）四元数分解为实部（w）和虚部（x, y, z）。

2. 计算标签四元数的共轭（label_inv）。

3. 将预测四元数（y_hat）分解为实部（sa）和虚部（xa, ya, za）。

4. 将标签共轭四元数（label_inv）分解为实部（sb）和虚部（xb, yb, zb）。

5. 计算两个四元数的乘积（loss_quat）。

6. 提取乘积四元数中的实部（cos_angle），它表示输入四元数之间的夹角的余弦值。

7. 计算损失值：1 减去余弦值的平方（loss）。

8. 返回损失值的均值。

这个损失函数的作用是衡量预测四元数（y_hat）与标签四元数（label）之间的角度差异。当预测四元数与标签四元数之间的角度相差较大时，损失值较大；当它们之间的角度相差较小时，损失值较小。这有助于神经网络在训练过程中学习到正确的旋转信息。

训练旋转四元数 TrainQuaternion

import argparse
import time
import torch
from torch.optim.lr_scheduler import StepLR

import DataLoader
import Net
import QuatLoss
import TransformLoss
import torch
from torch.utils.data import dataloader
import os

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=32, help='input batch size')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=0)
parser.add_argument('--nepoch', type=int, default=10000, help='number of epochs to train for')
parser.add_argument('--outf', type=str, default='seg', help='output folder')
parser.add_argument('--lr', type=float, default=0.1, help='learning rate')
parser.add_argument('--dataset', type=str, default='StandardModel', help="dataset path")
parser.add_argument('--npoints', type=int, default=500, help="sample point counts")
parser.add_argument('--stdModelPath', type=str, default='StandardModel', help="standard model path")

opt = parser.parse_args()

# 保存和输出设置
save_epoch = 100  # 几个epoch保存一次
println_epoch = 10  # 几个epoch输出一次loss


# 实例化网络
net = Net.RotationNet()
net.to(device)
# net = torch.load('./Save/1679852443_SGD_QuatLoss_batchsize32_/1679852443_epoch_4700_loss_0.01935810090508312.module')

# 优化器
optimizer = torch.optim.SGD(params=net.parameters(), lr=opt.lr, momentum=0.3)
optimizer_name = 'SGD'

# 损失函数
criterion = QuatLoss.QuatLoss()
criterion_name = 'QuatLoss'

# 学习率衰减
scheduler = StepLR(optimizer, step_size=1000, gamma=0.95)  # 学习率衰减

# DataLoader
train_dataset = DataLoader.LoadData(points_count=opt.npoints,
                                    std_model_paths=opt.stdModelPath,
                                    t_range=0.2,
                                    label_is_rotation=True)
# t_range 是 随机出的平移矩阵的范围， label_is_rotation 指数据的label是旋转四元数还是平移向量
train_loader = dataloader.DataLoader(dataset=train_dataset, batch_size=opt.batchSize, shuffle=True, num_workers=opt.workers)

# 模型保存路径
name = str(int(time.time()))
save_path = os.path.join('RotationNetSave', name + '_' + optimizer_name + '_' + criterion_name + '_batchsize' + str(opt.batchSize) + '_lr_' + str(opt.lr))
if not os.path.exists(save_path):
    os.makedirs(save_path)

# 记录损失值
loss_sum = 0
loss_count = 0

# 开始训练
for epoch in range(opt.nepoch):
    count = 0
    for points_align, points_std, label in train_loader:
        count += 1

        points_align = points_align.float().to(device)
        points_std = points_std.float().to(device)
        label = label.float().to(device)
        x1 = points_align.float().transpose(1, 2)
        x2 = points_std.float().transpose(1, 2)
        optimizer.zero_grad()
        y_hat = net(x1, x2).float()
        loss = criterion(y_hat, label)
        loss_sum += loss.item()
        loss_count += 1
        loss.backward()
        optimizer.step()

    scheduler.step()  # 更新学习率

    if epoch % save_epoch == 0:
        torch.save(net, save_path + "/" + name + '_epoch_' + str(epoch) + '_loss_' + str(loss_sum / loss_count) +'.module')

    if epoch % println_epoch == 0:
        print("epoch ", epoch, " loss ", loss_sum / loss_count, " lr ", get_lr(optimizer))
        loss_sum = 0
        loss_count = 0


torch.save(net, save_path + "/" + name + '.module')

1. 设置命令行参数以配置训练过程。

2. 初始化设备为GPU或CPU，取决于可用硬件。

3. 创建并配置平移网络模型、优化器、损失函数和学习率衰减策略。

4. 加载预训练的RotationNet模型，用于将点云旋转到正确的位姿。

5. 设置数据集和数据加载器。

6. 创建模型保存路径。

7. 进行训练循环，每个epoch中：

   • 遍历数据集的批次。

   • 将数据移动到指定设备（GPU或CPU）上。

   • 使用RotationNet网络调整输入点云的位姿。

   • 执行正向传播和反向传播。

   • 更新优化器和学习率。

   • 每隔一定步长保存一次模型

训练平移向量 TrainTranslateVector

import argparse
import time
import torch
from torch.optim.lr_scheduler import StepLR

import DataLoader
import Net
import QuatLoss
import torch
from torch.utils.data import dataloader
import os
import Quaternion as quat
import TransformLoss


def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=32, help='input batch size')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=0)
parser.add_argument('--nepoch', type=int, default=10000, help='number of epochs to train for')
parser.add_argument('--outf', type=str, default='seg', help='output folder')
parser.add_argument('--lr', type=float, default=0.1, help='learning rate')
parser.add_argument('--dataset', type=str, default='StandardModel', help="dataset path")
parser.add_argument('--npoints', type=int, default=500, help="sample point counts")
parser.add_argument('--stdModelPath', type=str, default='StandardModel', help="standard model path")

opt = parser.parse_args()

# 保存和输出设置
save_epoch = 100  # 几个epoch保存一次
println_epoch = 10  # 几个epoch输出一次loss

# 实例化平移网络
net = Net.TranslationNet()
net.to(device)

# 加载旋转网络
rotation_net = torch.load('./RotationNetSave/1680457393_SGD_QuatLoss_batchsize32_lr_0.1/1680457393_epoch_9900_loss_0.026239702384918927.module',
                          map_location=torch.device(device))

# 优化器
optimizer = torch.optim.SGD(params=net.parameters(), lr=opt.lr, momentum=0.3)
optimizer_name = 'SGD'

# 损失函数
# criterion = torch.nn.MSELoss()
# criterion_name = 'MSELoss'
criterion = TransformLoss.TransformLoss()
criterion_name = 'TransformLoss'

# 学习率衰减
scheduler = StepLR(optimizer, step_size=5000, gamma=0.95)  # 学习率衰减

# DataLoader
train_dataset = DataLoader.LoadData(points_count=opt.npoints,
                                    std_model_paths=opt.stdModelPath,
                                    t_range=0.2,
                                    label_is_rotation=False)
# t_range 是 随机出的平移矩阵的范围， label_is_rotation 指数据的label是旋转四元数还是平移向量
train_loader = dataloader.DataLoader(dataset=train_dataset,
                                     batch_size=opt.batchSize, shuffle=True,
                                     num_workers=opt.workers)

# 模型保存路径
name = str(int(time.time()))
save_path = os.path.join('TranslationNetSave', name + '_' + optimizer_name + '_' + criterion_name + '_batchsize' + str(opt.batchSize) + '_lr_' + str(opt.lr))
if not os.path.exists(save_path):
    os.makedirs(save_path)

# 记录损失值
loss_sum = 0
loss_count = 0

# 开始训练
for epoch in range(opt.nepoch):
    count = 0
    for points_align, points_std, label in train_loader:
        count += 1

        points_align = points_align.float().to(device)
        points_std = points_std.float().to(device)
        label = label.float().to(device)
        x1 = points_align.float().transpose(1, 2)
        x2 = points_std.float().transpose(1, 2)
        # 将x1和x2用RotationNet进行位姿调整
        r = rotation_net(x1, x2).float()  # 得到四元数
        batch_size = r.shape[0]  # 获取batchsize
        for i in range(batch_size):
            quat_obj = quat.Quaternion.from_tensor(r[i])  # 转四元数类
            R = quat_obj.getRotationMatrix().to(device)  # 得到R旋转矩阵
            x1[i] = torch.matmul(x1[i].transpose(0, 1), R).transpose(0, 1)  # 每一个点云都进行旋转

        optimizer.zero_grad()
        y_hat = net(x1, x2).float()
        loss = criterion(x1, x2, y_hat)
        loss_sum += loss.item()
        loss_count += 1
        loss.backward()
        optimizer.step()
        scheduler.step()  # 更新学习率

    if epoch % save_epoch == 0:
        torch.save(net, save_path + "/" + name + '_epoch_' + str(epoch) + '_loss_' + str(loss_sum / loss_count) +'.module')

    if epoch % println_epoch == 0:
        print("epoch ", epoch, " loss ", loss_sum / loss_count, " lr ", get_lr(optimizer))
        loss_sum = 0
        loss_count = 0


torch.save(net, save_path + "/" + name + '.module')

1. 导入所需库和模块，包括PyTorch、自定义的DataLoader、Net等。

2. 定义get_lr()函数，用于获取优化器的当前学习率。

3. 设定训练使用的设备（GPU或CPU）。

4. 定义命令行参数，包括批次大小、工作进程数、训练周期数、输出文件夹等。

5. 创建TranslationNet实例并将其加载到相应设备上。

6. 加载预训练的RotationNet模型。

7. 配置优化器、损失函数和学习率衰减策略。

8. 设置DataLoader来加载训练数据集。

9. 创建模型保存路径。

10. 初始化损失值记录变量。

11. 进行训练循环，每个循环执行以下操作：

    • 遍历数据集的批次。

    • 将数据移动到指定设备（GPU或CPU）。

    • 使用RotationNet网络调整点云的位姿。

    • 执行正向传播、计算损失、执行反向传播。

    • 更新优化器和学习率。

    • 每隔一定周期保存模型和输出损失值。

测试代码 Test

import open3d as o3d
import numpy as np
import torch
from torch.utils.data import dataloader

import DataLoader
import QuatLoss
import Quaternion as quat
import ICP

number_of_points = 500
rotation_net = torch.load('./RotationNetSave/1680457393_SGD_QuatLoss_batchsize32_lr_0.1/1680457393_epoch_9900_loss_0.026239702384918927.module',
                 map_location=torch.device('cpu'))
transform_net = torch.load('./TranslationNetSave/1680461748_SGD_MSELoss_batchsize32_lr_0.1/1680461748_epoch_13000_loss_0.0005985869252072007.module',
                 map_location=torch.device('cpu'))

test_dataset = DataLoader.LoadData(points_count=500, std_model_paths='StandardModel', t_range=0.2, label_is_rotation=False)
test_loader = dataloader.DataLoader(dataset=test_dataset, batch_size=32, shuffle=False)

criterion = QuatLoss.QuatLoss()
# 是否开启ICP的细配准
use_icp = True

for points_align, points_std, label in test_loader:
    x1 = points_align.float().transpose(1, 2)
    x2 = points_std.float().transpose(1, 2)
    y_hat_rotation = rotation_net(x1, x2).float()
    # loss = criterion(y_hat_rotation, label)
    x1_single = x1[0].transpose(0, 1)
    x2_single = x2[0].transpose(0, 1)
    # 输入 需对其点云
    pcd_before = o3d.geometry.PointCloud()
    pcd_before.points = o3d.utility.Vector3dVector(x1_single)
    pcd_before.paint_uniform_color([0, 1, 0])  # 指定显示为绿色

    # 将输入点云进行预测值的旋转
    r = rotation_net(x1, x2)
    batch_size = r.shape[0]
    for i in range(batch_size):
        quat_obj = quat.Quaternion.from_tensor(r[i])
        R = quat_obj.getRotationMatrix()
        x1[i] = torch.matmul(x1[i].transpose(0, 1), R).transpose(0, 1)

    # 进行平移
    y_hat_transform = transform_net(x1, x2).float()
    x1 = x1.transpose(1, 2)
    x1 = x1 + y_hat_transform.unsqueeze(1)
    x2 = x2.transpose(1, 2)

    points_after = x1[0].detach()
    if use_icp:
        # 用ICP算法进行细调整
        R, t = ICP.ICP(x1[0].detach().numpy(), x2[0].detach().numpy(), 80)
        points_after_ICP = torch.matmul(x1[0], torch.from_numpy(R).float())
        t = torch.from_numpy(t.reshape(1, 3)).float()
        points_after_ICP += t

    # 显示旋转后的模型
    pcd_after = o3d.geometry.PointCloud()
    pcd_after.points = o3d.utility.Vector3dVector(points_after.detach())
    pcd_after.paint_uniform_color([1, 0, 0])  # 红色

    # 显示std模型
    pcd_std = o3d.geometry.PointCloud()
    pcd_std.points = o3d.utility.Vector3dVector(x2_single)
    pcd_std.paint_uniform_color([0, 0, 0])  # 黑色

    if use_icp:
        pcd_after.points = o3d.utility.Vector3dVector(points_after_ICP.detach())
        o3d.visualization.draw_geometries([pcd_before, pcd_after, pcd_std])  # 显示对其后位置
    else:
        o3d.visualization.draw_geometries([pcd_before, pcd_after, pcd_std])  # 显示对其后位置
ython

这段代码实现了使用训练好的RotationNet和TranslationNet进行点云对齐。主要步骤如下：

1. 导入所需库，包括Open3D、NumPy、PyTorch等。

2. 加载训练好的RotationNet和TranslationNet模型。

3. 设置测试数据集和DataLoader。

4. 定义四元数损失函数。

5. 设置是否使用ICP算法（迭代最近点算法）进行细对齐。

6. 遍历测试数据集的批次，执行以下操作：

   • 将输入点云和标准点云传入RotationNet，得到预测的旋转四元数。

   • 使用预测的旋转四元数对输入点云进行旋转。

   • 将旋转后的点云和标准点云传入TranslationNet，得到预测的平移向量。

   • 使用预测的平移向量对输入点云进行平移。

   • 如果使用ICP算法，则对旋转和平移后的点云进行细对齐。

   • 将输入点云、对齐后的点云和标准点云可视化并显示。

在这个脚本中，使用Open3D库进行点云的可视化和处理。通过预测的旋转四元数和平移向量，对输入点云进行旋转和平移，以实现与标准点云的对齐。如果启用了ICP算法，还会对旋转和平移后的点云进行细对齐。最后，将输入点云、对齐后的点云和标准点云可视化并显示。