【python深度学习】——torch.expand广播机制|torch.norm
- 1. torch.expand()与广播机制
- 2. torch.norm()计算范数
- 3. 结合使用的完整示例代码
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 npdef 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_expandeddef 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 cloudpoint_cloud = np.transpose((R @ np.array(point_cloud).T) + np.tile(T,(1, point_cloud.shape[0])))return point_cloudif __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}")