torch.randint 用于生成指定范围内的随机整数张量，是常用的张量创建方法之一。它特别适用于需要生成离散随机数的时候，如初始化张量、创建随机索引等。

方法签名

torch.randint(low=0, high, size, out=None, dtype=torch.int64, layout=torch.strided, device=None, requires_grad=False)

参数详解

  • low (int, optional)：生成随机整数的下界（包含），默认为0。如果不设置，则默认最低值为0。
  • high (int)：生成随机整数的上界（不包含）。
  • size (tuple of ints)：新生成的张量的形状。
  • out (Tensor, optional)：输出张量，如果指定，结果将被存储到该张量中。
  • dtype (torch.dtype, optional)：输出张量的数据类型，默认是 torch.int64。
  • layout (torch.layout, optional)：张量的布局，默认是 torch.strided。
  • device (torch.device, optional)：张量所在的设备（CPU或CUDA），默认情况下根据运行环境自动选择。
  • requires_grad (bool, optional)：如果设置为True，则创建的张量会记录其计算图，用于自动求导，默认是False。

返回值

返回一个包含在给定范围内随机整数的张量，其形状由size参数指定。

例子

基本用法

import torch 
# 生成一个形状为 [5] 的张量，包含5个在 0 到 10 之间的随机整数 
tensor1 = torch.randint(0, 10, (5,)) print(tensor1)

指定形状的张量

import torch

# 生成一个形状为 [2, 3] 的张量，包含在 -5 到 5 之间的随机整数
tensor2 = torch.randint(-5, 5, (2, 3))
print(tensor2)

指定设备的张量

import torch

# 生成一个形状为 [2, 3] 的张量，包含在 0 到 10 之间的随机整数，并指定生成的张量在CUDA设备上
tensor3 = torch.randint(0, 10, (2, 3), device='cuda')
print(tensor3)
print(tensor3.device)

生成requires_grad的张量

import torch

# 生成一个形状为 [2, 2] 的张量，包含在 0 到 5 之间的随机整数，并且该张量记录计算图以用于自动求导
tensor4 = torch.randint(0, 5, (2, 2), requires_grad=True)
print(tensor4)
print(tensor4.requires_grad)

使用out参数

import torch

# 预先定义一个张量
out_tensor = torch.empty(3, 3)

# 使用 randint 方法将随机整数填入预先定义的张量
torch.randint(0, 10, (3, 3), out=out_tensor)
print(out_tensor)

注意事项

  • high 参数必须指定，且必须大于 low 参数。
  • 如需生成在最小值和最大值之间的张量，请确保 high 和 low 的值设置正确。
  • 如果需要在某一特定设备（如GPU）上创建张量，可以通过设置 device 参数来实现。
  • 当需要记录计算图以便于梯度计算时，可以设置 requires_grad=True。

torch.randint 是生成离散均匀分布随机数张量的常用工具，它的灵活性和易用性使得这个方法在许多场景中都非常实用。


Pytorch的torch.randint()详解插图

关注公众号：程序新视界，一个让你软实力、硬技术同步提升的平台

除非注明，否则均为程序新视界原创文章，转载必须以链接形式标明本文链接

本文链接：https://choupangxia.com/2024/10/04/pytorch-torch-randint/