【PyTorch】torch.randint的详细用法

目录

torch.randint的详细用法

语法：

参数：

返回值：

示例 1：生成一个 2x3 的随机整数张量

输出：

示例 2：生成指定数据类型的随机整数张量

输出：

示例 3：生成张量并指定设备

示例 4：生成 1D 张量

输出：

示例 5：生成带有 requires_grad=True 的张量

输出：

示例 6：生成多维张量

输出：

总结：

---

torch.randint的详细用法

torch.randint() 是 PyTorch 中用于生成随机整数张量的函数。

它会返回一个包含指定形状的张量，张量中的每个元素是从指定的整数区间内随机选取的整数。

语法：

torch.randint(low, high, size, dtype=None, device=None, requires_grad=False)

参数：

• low：整数类型，生成随机整数的下限（包括该值）。
• high：整数类型，生成随机整数的上限（不包括该值）。low 必须小于 high。
• size：形状（tuple 或 list），指定返回张量的形状，例如 (2, 3) 表示一个 2x3 的张量。
• dtype（可选）：张量的数据类型，默认是 torch.int64，可以指定为 torch.int32 或其他整数类型。
• device（可选）：指定张量生成在哪个设备上（如 cpu 或 cuda）。默认为 None，表示使用默认设备（通常是 CPU）。
• requires_grad（可选）：如果设置为 True，则会记录操作的梯度，用于自动求导。默认值是 False。

返回值：

返回一个整数类型的张量，其形状由 size 参数指定，元素是从 [low, high) 区间内随机选择的整数。

---

示例 1：生成一个 2x3 的随机整数张量

import torch

# 生成一个 2x3 的随机整数张量，元素值在 [0, 10) 之间
random_tensor = torch.randint(0, 10, (2, 3))
print(random_tensor)

输出：

tensor([[8, 1, 4],
        [7, 2, 3]])

解释：

• 生成了一个形状为 (2, 3) 的随机张量，元素值在 [0, 10) 之间，表示随机生成的整数是 [0, 1, ..., 9] 中的数。

---

示例 2：生成指定数据类型的随机整数张量

# 创建一个数据类型为 torch.int32 的随机整数张量
random_tensor_int32 = torch.randint(0, 100, (2, 2), dtype=torch.int32)
print(random_tensor_int32)

输出：

tensor([[87, 34],
        [65, 22]], dtype=torch.int32)

解释：

• 这里生成的张量数据类型为 torch.int32，其元素值是 [0, 100) 之间的随机整数。

---

示例 3：生成张量并指定设备

# 如果 CUDA 可用，将张量生成在 GPU 上
if torch.cuda.is_available():
    random_tensor_gpu = torch.randint(0, 100, (2, 3), device='cuda')
    print(f"Random Tensor on GPU: {random_tensor_gpu}")
else:
    print("CUDA is not available.")

解释：

• 这个示例中，如果你的机器有可用的 GPU，张量 random_tensor_gpu 会生成在 GPU 上。否则，它会打印 "CUDA is not available."。

---

示例 4：生成 1D 张量

# 生成一个 1D 随机整数张量，形状为 (5,)
random_tensor_1d = torch.randint(0, 10, (5,))
print(random_tensor_1d)

输出：

tensor([8, 4, 2, 9, 3])

解释：

• 生成了一个包含 5 个元素的一维张量，元素值在 [0, 10) 之间。

---

示例 5：生成带有 requires_grad=True 的张量

# 生成一个随机整数张量，要求计算梯度
random_tensor_grad = torch.randint(0, 10, (2, 3), requires_grad=True)
print(random_tensor_grad)

输出：

tensor([[6, 2, 7],
        [2, 1, 9]], requires_grad=True)

解释：

• requires_grad=True 表示生成的张量会记录操作，并参与反向传播计算梯度。
• 如果你对该张量进行进一步操作（如加法、乘法等），PyTorch 会记录这些操作的计算图，便于后续进行梯度计算。

---

示例 6：生成多维张量

# 生成一个 3 维张量，形状为 (2, 3, 4)
random_tensor_3d = torch.randint(0, 100, (2, 3, 4))
print(random_tensor_3d)

输出：

tensor([[[ 2,  9, 49, 69],
         [71, 77, 51, 58],
         [95, 38,  0, 29]],

        [[52, 73, 56, 91],
         [18, 38, 34, 12],
         [29, 18, 64, 75]]])

解释：

• 这里生成了一个形状为 (2, 3, 4) 的三维张量，表示有 2 个样本，每个样本包含 3 行 4 列的随机整数。

---

总结：

• torch.randint() 用于生成指定形状的随机整数张量，其元素从区间 [low, high) 中随机选择。
• 它的常见用法包括生成不同形状的张量，指定数据类型、设备（如 CPU 或 GPU）以及是否需要求导。
• 生成的随机整数张量的元素值总是落在 [low, high) 范围内，low 是包含的，high 是不包含的。