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()详解插图

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

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

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