现在位置: 首页 > PyTorch 教程 > 正文

PyTorch torch.arange 函数


Pytorch torch 参考手册 Pytorch torch 参考手册

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)

输出结果为:

tensor([0, 1, 2, 3, 4])

示例 2: 指定起始值和结束值

实例

import torch

# 创建 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)

输出结果为:

tensor([0, 2, 4, 6, 8, 10])

示例 4: 负数步长

实例

import torch

# 创建 10 到 0,步长为 -2 的等差数列
x = torch.arange(10, 0, -2)

print(x)

输出结果为:

tensor([10,  8,  6,  4,  2])

Pytorch torch 参考手册 Pytorch torch 参考手册