介绍
生成连续整数/浮点数序列张量的核心工具，常用于创建位置编码、索引、数据生成等场景（比如你之前关注的大模型位置信息注入就会用到它）。
功能
torch.arange() 生成一个一维张量，包含从start到end（不包含end）、步长为step的序列，类似Python的range()，但返回的是PyTorch张量（支持GPU加速）。
基本语法
torch.arange(start=0, end, step=1, dtype=None, device=None, requires_grad=False)
| 参数 | 作用 |
|---|---|
start |
序列起始值（默认0） |
end |
序列结束值（不包含，必填） |
step |
步长（默认1，可正可负，不能为0） |
dtype |
张量数据类型（如torch.int32/torch.float32，默认自动推断） |
device |
张量存储设备（如cuda:0/cpu） |
requires_grad |
是否需要计算梯度（默认False） |
例子
1. 基础用法（默认参数）
import torch
# 生成 0~4 的整数序列（不包含5）
a = torch.arange(5)
print(a) # 输出：tensor([0, 1, 2, 3, 4])
print(a.dtype) # 输出：torch.int64（默认整数类型）
2. 指定起始值和步长
# 生成 2~10，步长2
b = torch.arange(2, 10, 2)
print(b) # 输出：tensor([2, 4, 6, 8])
# 生成 10~2，步长-2（倒序）
c = torch.arange(10, 2, -2)
print(c) # 输出：tensor([10, 8, 6, 4])
3. 生成浮点数序列
# 指定dtype为float，生成 0~1，步长0.2
d = torch.arange(0, 1, 0.2, dtype=torch.float32)
print(d) # 输出：tensor([0.0000, 0.2000, 0.4000, 0.6000, 0.8000])
4. 结合设备（GPU）
# 生成序列并直接放到GPU（需有CUDA环境）
if torch.cuda.is_available():
e = torch.arange(5, device="cuda:0")
print(e) # 输出：tensor([0, 1, 2, 3, 4], device='cuda:0')
应用场景
1. 生成位置编码（大模型核心用法）
这是你之前关注的“注入位置信息”的关键步骤，用torch.arange()生成位置索引：
# 生成序列长度为10的位置索引（0~9）
seq_len = 10
position = torch.arange(0, seq_len, dtype=torch.float32)
print(position) # 输出：tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
# 基于位置索引生成正弦位置编码（Transformer原版）
d_model = 16 # 嵌入维度
div_term = torch.exp(torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model))
pe = torch.zeros(seq_len, d_model)
pe[:, 0::2] = torch.sin(position.unsqueeze(1) * div_term) # 偶数维度
pe[:, 1::2] = torch.cos(position.unsqueeze(1) * div_term) # 奇数维度
print("位置编码形状：", pe.shape) # 输出：torch.Size([10, 16])
2. 生成索引/切片
# 生成索引，用于提取张量的前5个元素
x = torch.randn(10, 3) # 10行3列的随机张量
idx = torch.arange(5)
x_slice = x[idx] # 提取前5行
print(x_slice.shape) # 输出：torch.Size([5, 3])
注意事项（避坑）
end是开区间：torch.arange(0, 5)生成[0,1,2,3,4]，不包含5；- 步长不能为0：会报
ValueError； - 数据类型：整数步长默认
int64，包含浮点数步长（如0.2）自动转为float32； - 与
torch.linspace()的区别：arange：按步长生成（不保证点数）；linspace：按点数生成（均分区间，包含end）。
总结
torch.arange()是PyTorch生成连续序列张量的核心函数，语法类似Pythonrange()但支持张量特性；- 核心参数：
start（起始）、end（结束，不包含）、step（步长）； - 高频场景：生成大模型位置编码、张量索引、数据序列生成，是LLM开发中注入位置信息的基础工具。