PyTorch中分段曲线段长度计算：是否有更优实现方式？

结论 #

你当前的循环实现并非最优，没有利用PyTorch核心的向量化并行计算能力，当曲线点数量较大时效率会明显偏低。下面是更符合PyTorch风格的实现方式：

方法一：直接用张量切片+范数计算（最简洁高效） #

这是最优的实现，完全基于张量操作，避免任何Python循环：

import torch

line = torch.tensor([(-104.6400,0.0000),(-104.6400,0.1500),(-103.5500,0.5140),(-98.1000,1.0775),(-92.6500,1.4553)])
# 计算相邻点的差值张量
diff = line[1:] - line[:-1]
# 对每个差值向量计算L2范数（欧氏距离），得到各段长度
segment_lengths = torch.norm(diff, dim=1, p=2)

# 按索引输出结果
for idx, length in enumerate(segment_lengths, start=1):
    print(f"{idx}: {length}")

方法二：用cdist实现（适合需要批量计算点对距离的场景） #

如果需要基于cdistAPI实现，同样可以向量化处理，无需循环：

import torch
from torch.nn.functional import cdist

line = torch.tensor([(-104.6400,0.0000),(-104.6400,0.1500),(-103.5500,0.5140),(-98.1000,1.0775),(-92.6500,1.4553)])
# 提取相邻点对的两组张量
prev_points = line[:-1]
curr_points = line[1:]
# 计算点对距离矩阵，取对角线元素（对应相邻点的距离）
segment_lengths = cdist(prev_points, curr_points, p=2).diag()

# 输出结果
for idx, length in enumerate(segment_lengths, start=1):
    print(f"{idx}: {length}")

为什么这些方法更好？ #

• 效率更高：向量化操作能利用PyTorch的CPU/GPU并行计算能力，当点数量达到数千甚至上万时，速度会比循环实现快几十倍。
• 代码更简洁：避免手动维护循环索引、中间变量，符合PyTorch的张量编程范式。
• 扩展性强：如果需要计算其他类型的距离（比如曼哈顿距离p=1），只需修改norm或cdist的p参数即可。

内容的提问来源于stack exchange，提问作者FlumeRS