pytorch api文档：torch.arange()-马育民老师

# 介绍

生成**连续整数/浮点数序列张量**的核心工具，常用于创建位置编码、索引、数据生成等场景（比如你之前关注的大模型位置信息注入就会用到它）。

# 功能
`torch.arange()` 生成一个**一维张量**，包含从`start`到`end`（不包含`end`）、步长为`step`的序列，类似Python的`range()`，但返回的是PyTorch张量（支持GPU加速）。

### 基本语法
```python
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. 基础用法（默认参数）
```python
import torch

# 生成 0~4 的整数序列（不包含5）
a = torch.arange(5)
print(a)  # 输出：tensor([0, 1, 2, 3, 4])
print(a.dtype)  # 输出：torch.int64（默认整数类型）
```

### 2. 指定起始值和步长
```python
# 生成 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. 生成浮点数序列
```python
# 指定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）
```python
# 生成序列并直接放到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()`生成位置索引：
```python
# 生成序列长度为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. 生成索引/切片
```python
# 生成索引，用于提取张量的前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])
```

# 注意事项（避坑）

1. `end`是**开区间**：`torch.arange(0, 5)`生成`[0,1,2,3,4]`，不包含5；
2. 步长不能为0：会报`ValueError`；
3. 数据类型：整数步长默认`int64`，包含浮点数步长（如0.2）自动转为`float32`；
4. 与`torch.linspace()`的区别：
   - `arange`：按**步长**生成（不保证点数）；
   - `linspace`：按**点数**生成（均分区间，包含end）。

# 总结
1. `torch.arange()`是PyTorch生成连续序列张量的核心函数，语法类似Python`range()`但支持张量特性；
2. 核心参数：`start`（起始）、`end`（结束，不包含）、`step`（步长）；
3. 高频场景：生成大模型位置编码、张量索引、数据序列生成，是LLM开发中注入位置信息的基础工具。

原文出处：http://www.malaoshi.top/show_1GW2YuyIR0Xi.html