介绍

torch.randint() 函数，是生成 整数型随机张量 的核心工具，能指定整数的取值范围和张量形状，常用来生成标签、索引、随机掩码等离散型随机数据

作用

创建指定形状的张量，其中的元素值随机采样自 [low, high) 区间的整数（包含 low，不包含 high），且每个整数被采样的概率相等（均匀分布）。

是唯一专门生成整数随机张量的基础函数，对比浮点随机函数：

• torch.rand()：[0,1) 浮点均匀分布；
• torch.randn()：标准正态分布浮点；

语法

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

参数

例子

high 和 size 是必传参数（low 可选，默认 0），直接看代码更直观：

import torch

# 示例1：基础用法（low默认0，生成 [0,5) 的整数，形状[3]）
randint_1d = torch.randint(5, size=(3,))
print("一维整数张量：", randint_1d)
# 输出示例：tensor([2, 0, 4])（值∈{0,1,2,3,4}）
print("形状：", randint_1d.shape)  # torch.Size([3])

# 示例2：指定low和high（生成 [2,7) 的整数，形状[2,3]）
randint_2d = torch.randint(low=2, high=7, size=(2, 3))
print("二维整数张量：\n", randint_2d)
# 输出示例：
# tensor([[4, 2, 6],
#         [3, 5, 2]])（值∈{2,3,4,5,6}）

# 示例3：高维张量（批量标签，形状[4,2]，生成 [0,10) 的整数）
randint_3d = torch.randint(0, 10, size=(4, 2))
print("批量标签张量：\n", randint_3d)
# 输出示例：
# tensor([[7, 9],
#         [1, 4],
#         [8, 0],
#         [3, 5]])

参数示例

# 示例1：指定数据类型（int32）
randint_int32 = torch.randint(0, 5, size=(2, 3), dtype=torch.int32)
print("数据类型：", randint_int32.dtype)  # torch.int32
print("数值：\n", randint_int32)

# 示例2：指定GPU设备（需CUDA环境）
if torch.cuda.is_available():
    randint_cuda = torch.randint(0, 10, size=(2, 3), device="cuda")
    print("设备：", randint_cuda.device)  # cuda:0
else:
    print("无CUDA环境，默认生成CPU张量")

# 示例3：low=high-1（所有元素都是同一个值）
randint_single = torch.randint(3, 4, size=(3, 3))
print("单值整数张量：\n", randint_single)  # 全为3

随机种子（复现结果）

和其他随机函数一样，设置全局种子可固定 torch.randint() 的结果，保证实验可复现：

# 设置全局随机种子
torch.manual_seed(42)
print("固定种子结果1：", torch.randint(0, 10, size=(2,)))  # tensor([2, 7])
torch.manual_seed(42)
print("固定种子结果2：", torch.randint(0, 10, size=(2,)))  # tensor([2, 7])（完全一致）

与 torch.rand() 的转换关系

如果需要将 torch.rand() 的浮点结果转为整数，可通过缩放+取整实现，但 torch.randint() 更直接：

# 需求：生成 [0,5) 的整数（两种方式）
# 方式1：torch.randint（推荐，直接生成）
res1 = torch.randint(0, 5, size=(3,))
# 方式2：torch.rand 缩放后取整
res2 = (torch.rand(3) * 5).int()

print("randint 结果：", res1)
print("rand 转换结果：", res2)
# 两种方式结果等价，但 randint 更高效、语义更清晰

注意

整数张量的限制

torch.randint() 生成的整数张量无法直接参与梯度计算，若需作为可训练参数，需先转换为浮点型：

# 错误示例：整数张量设 requires_grad=True 无意义
randint_grad = torch.randint(0, 5, size=(2,3), requires_grad=True)
print("是否可求梯度：", randint_grad.requires_grad)  # True（但实际无法计算梯度）

# 正确：转换为浮点型后再开启梯度
randint_float = torch.randint(0, 5, size=(2,3)).float()
randint_float.requires_grad = True
print("浮点张量是否可求梯度：", randint_float.requires_grad)  # True

实战场景

生成分类任务标签

这是最常用的场景，生成指定类别数的随机标签：

# 场景：生成100个样本的分类标签（类别数：0~9）
batch_size = 100
num_classes = 10
labels = torch.randint(0, num_classes, size=(batch_size,))

print("标签形状：", labels.shape)  # torch.Size([100])
print("标签示例：", labels[:10])  # tensor([5, 3, 8, 1, 9, 0, 7, 2, 4, 6])
print("标签范围：", (labels.min().item(), labels.max().item()))  # (0,9)

生成随机索引（数据采样/打乱）

用于随机选取数据、打乱序列等场景：

# 场景：从100个样本中随机选取10个（生成随机索引）
total_samples = 100
select_num = 10
indices = torch.randint(0, total_samples, size=(select_num,))

print("随机索引：", indices)  # tensor([45, 12, 89, 3, 77, 21, 90, 5, 63, 18])

# 模拟采样（用索引选取数据）
data = torch.randn(total_samples, 10)  # 100个样本，10维特征
selected_data = data[indices]
print("采样后数据形状：", selected_data.shape)  # torch.Size([10, 10])

生成随机掩码（离散值）

生成 0/1 掩码（如 dropout 掩码、注意力掩码）：

# 场景：生成4×10的0/1掩码（50%概率为1）
# 方式：生成 [0,2) 的整数（0或1）
mask = torch.randint(0, 2, size=(4, 10))
print("随机0/1掩码：\n", mask)
# 输出示例：
# tensor([[1, 0, 1, 0, 1, 0, 1, 0, 1, 0],
#         [0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
#         [1, 1, 0, 0, 1, 1, 0, 0, 1, 1],
#         [0, 0, 1, 1, 0, 0, 1, 1, 0, 0]])

总结

1. torch.randint(low, high, size) 生成 [low, high) 区间的整数随机张量，是PyTorch中唯一的整数随机生成函数。
2. 核心参数：high（上界，必传）、size（形状，需传元组），low 默认0，dtype 默认int64。
3. 关键限制：整数张量无法直接求梯度，需转换为浮点型；size 参数必须用元组包裹（区别于rand/randn）。
4. 适用场景：生成分类标签、随机索引、离散掩码等离散型随机数据，浮点随机场景用rand/randn。

原文出处：/show_1GW2c5w3kyYY.html