【python深度学习】——torch.expand()广播机制|torch.norm()

1. torch.expand()与广播机制

在处理3D点云时, 有时需要对两帧点云进行逐点的三维坐标相加减、做点积等运算, 但是读入的PCD文件中,点云数量并不一定是相等的

那么首要的一个问题就是, 如何将两帧点云处理成大小相同的矩阵然后进行计算?

torch 中一个常用的方法是expand函数, 例如下面这段函数:

def expand_matrix(ori_A, ori_B):
    point_num_A = ori_A.size()[0] #原始点云A的点数
    point_dim = ori_A.size()[1]   #原始点云A的坐标维度,一般为3维
    point_num_B = ori_B.size()[0]
    assert point_dim == ori_B.size()[1] #保障都是3维点云

    # 假定expand目标为（point_num_A, point_num_B, point_dim）

    # 对A进行广播
    A_expanded = ori_A.unsqueeze(1)
    print(f"expand A,shape= {A_expanded.shape}, target = {point_num_A, point_num_B, point_dim}")
    A_expanded = A_expanded.expand(point_num_A, point_num_B, point_dim)

    # 对B进行广播
    B_expanded = ori_B.unsqueeze(0)
    print(f"expand B,shape= {B_expanded.shape}, target = {point_num_A, point_num_B, point_dim}")
    B_expanded = B_expanded.expand(point_num_A, point_num_B, point_dim)
    return A_expanded, B_expanded

代码解释:

expand函数只能在原始 维度为1 的情况下进行扩展。如果尝试在一个不是 1 的维度上进行扩展，会引发错误。因此我们先用unsqueeze函数在目标位置上进行dim=1的扩充, 再进行expand。

注意事项

• expand 方法不会实际分配新的内存，而是通过引用的方式实现扩展。
• 使用 expand 后的张量仍然与原始张量共享底层数据。

2. torch.norm()计算范数

torch.norm()是 PyTorch 中用于计算张量范数的函数,
基本用法如下：

torch.norm(input, p='fro', dim=None, keepdim=False, out=None, dtype=None)

这里的参数说明:

• input: 输入张量。
• p: 范数类型，默认为 ‘fro’（弗罗贝尼乌斯范数）。可以是以下值之一：
  ‘fro’ 或 None：弗罗贝尼乌斯范数（适用于矩阵）。
  2：默认的 L2 范数（适用于向量）。
  1：L1 范数。
  float(‘inf’) 或 ‘inf’：无穷范数。
  -float(‘inf’) 或 ‘-inf’：负无穷范数。
  其他正整数或浮点数：计算相应的 p-范数。
• dim: 要计算范数的维度。如果为 None，则计算整个张量的范数。
• keepdim: 是否保持输出张量的维度。如果为 True，则结果会保留被计算范数的维度。
• out: 输出张量。
• dtype: 返回张量的期望数据类型。如果为 None，则与输入数据类型相同。

3. 结合使用的完整示例代码

import torch
import open3d as o3d
import numpy as np

def expand_matrix(ori_A, ori_B):
    point_num_A = ori_A.size()[0]
    point_dim = ori_A.size()[1]    
    point_num_B = ori_B.size()[0]
    assert point_dim == ori_B.size()[1]

    # expand目标为（point_num_A, point_num_B, point_dim）

    # 对A进行广播
    A_expanded = ori_A.unsqueeze(1)
    print(f"expand A,shape= {A_expanded.shape}, target = {point_num_A, point_num_B, point_dim}")
    A_expanded = A_expanded.expand(point_num_A, point_num_B, point_dim)

    # 对B进行广播
    B_expanded = ori_B.unsqueeze(0)
    print(f"expand B,shape= {B_expanded.shape}, target = {point_num_A, point_num_B, point_dim}")
    B_expanded = B_expanded.expand(point_num_A, point_num_B, point_dim)
    return A_expanded, B_expanded

def calculate_distance(point_cloud_A, point_cloud_B):

    point_cloud_A = torch.tensor(point_cloud_A, dtype=torch.float32)
    point_cloud_B = torch.tensor(point_cloud_B, dtype=torch.float32)
    point_cloud_A = point_cloud_A.cuda()
    point_cloud_B = point_cloud_B.cuda()
    print("输入的A shape = ",point_cloud_A.shape)
    print("输入的B shape = ",point_cloud_B.shape)

    # 对A和B分别进行广播
    A_expanded, B_expanded = expand_matrix(point_cloud_A, point_cloud_B)

    # 计算每个点到另一个点云中所有点的距离, 结果的shape是(M,N)
    diff = torch.norm(A_expanded - B_expanded, dim=2, keepdim=False)
    for i in range(point_cloud_A.shape[0]):
        print(f"A<{i}>到B集合中所有点的距离: ",diff[i])
    for i in range(point_cloud_B.shape[0]):
        print(f"B<{i}>到A集合中所有点的距离: ",diff[:,i])


# 读取pcd点云
def read_pcd(file_path, main_pcd=False):
    pcd_data = o3d.io.read_point_cloud(file_path)
    # 进行下采样,避免点数过多，显存不够
    pcd_data = pcd_data.voxel_down_sample(voxel_size=1)
    point_cloud = np.asarray(pcd_data.points)
    if not main_pcd:
        transform = np.loadtxt('data/ndt.txt')
        R = transform[:3, :3]
        T = transform[:3, 3]
        T= np.expand_dims(T, axis=1) 
        # print("==== R is ", R)
        # print("==== T is ", T)
        # apply to point cloud
        point_cloud = np.transpose((R @ np.array(point_cloud).T) + np.tile(T,(1, point_cloud.shape[0])))
    return point_cloud

if __name__ == "__main__":

    ##! 读取目标点云
    # tgt = read_pcd('data/left_nonground.pcd',main_pcd=True)
    ##! 读取待转换的点云，并且进行RT变换
    # src = read_pcd('data/right_nonground.pcd')

    ##! 也可以使用测试数据,更为直观 
    tgt =  np.array([[1, 2, 3], [4,5,6], [1, 2, 3], [4,5,6],[1, 2, 3], [4,5,6] ])
    src = np.array([[0, 0,0],[1,1,1],[2,2,2],[3,3,3]])

    chamfer_dist_A, chamfer_dist_B = calculate_distance(tgt, src)
    print(f"Chamfer Distance from A to B: {chamfer_dist_A}")
    print(f"Chamfer Distance from B to A: {chamfer_dist_B}")