PyTorch torch.arange 函数
torch.arange 是 PyTorch 中用于创建等差数列张量的函数。它会创建一个一维张量,包含从起始值到结束值的等差数列。
这在深度学习中常用于创建索引、范围数组或在循环中生成序列。
函数定义
torch.arange(start=0, end, step=1, dtype=None, device=None, requires_grad=False)
参数:
start(float, 可选): 序列的起始值,默认为 0。end(float): 序列的结束值(不包含)。step(float, 可选): 步长,默认为 1。dtype(torch.dtype, 可选): 指定张量的数据类型。device(torch.device, 可选): 指定张量存储的设备。requires_grad(bool, 可选): 是否需要计算梯度。
返回值:
torch.Tensor: 返回一个一维张量。
使用示例
示例 1: 从 0 到 5
实例
import torch
# 创建 0 到 4 的等差数列
x = torch.arange(5)
print(x)
# 创建 0 到 4 的等差数列
x = torch.arange(5)
print(x)
输出结果为:
tensor([0, 1, 2, 3, 4])
示例 2: 指定起始值和结束值
实例
import torch
# 创建 2 到 8 的等差数列
x = torch.arange(2, 9)
print(x)
# 创建 2 到 8 的等差数列
x = torch.arange(2, 9)
print(x)
输出结果为:
tensor([2, 3, 4, 5, 6, 7, 8])
示例 3: 指定步长
实例
import torch
# 创建 0 到 10,步长为 2 的等差数列
x = torch.arange(0, 11, 2)
print(x)
# 创建 0 到 10,步长为 2 的等差数列
x = torch.arange(0, 11, 2)
print(x)
输出结果为:
tensor([0, 2, 4, 6, 8, 10])
示例 4: 负数步长
实例
import torch
# 创建 10 到 0,步长为 -2 的等差数列
x = torch.arange(10, 0, -2)
print(x)
# 创建 10 到 0,步长为 -2 的等差数列
x = torch.arange(10, 0, -2)
print(x)
输出结果为:
tensor([10, 8, 6, 4, 2])

Pytorch torch 参考手册