介绍
torch.ones() 函数,生成全 1 张量的基础工具,在张量初始化场景中高频使用。
与 torch.ones() 对应的是 torch.zeros()(生成全 0 张量),二者用法完全一致,只是填充值不同。
语法
生成一个指定形状、所有元素均为 1 的张量。
torch.ones(*size, dtype=None, device=None, requires_grad=False) → Tensor
参数:
*size:必选参数,指定张量的形状,可以是单个整数、多个整数,或元组/列表形式(如4、4,4、(4,4));dtype:可选,指定张量数据类型(如torch.float32、torch.int64),默认与 PyTorch 全局默认类型一致(通常是torch.float32);device:可选,指定张量存储设备(如cpu、cuda:0),默认在 CPU 上;requires_grad:可选,是否需要计算梯度,默认False(非训练参数无需开启)。
例子
import torch
# 示例1:创建基础的全1张量(指定形状为单个整数)
# 创建一个长度为5的一维全1张量
tensor1 = torch.ones(5)
print("示例1 - 一维全1张量:")
print(tensor1) # 输出:tensor([1., 1., 1., 1., 1.])
# 示例2:创建多维全1张量(指定形状为多个整数)
# 创建 2行3列 的二维全1张量
tensor2 = torch.ones(2, 3)
print("\n示例2 - 二维全1张量:")
print(tensor2)
# 输出:
# tensor([[1., 1., 1.],
# [1., 1., 1.]])
# 示例3:指定数据类型和设备
# 创建 3行2列、数据类型为int64、存储在CPU上的全1张量
tensor3 = torch.ones((3, 2), dtype=torch.int64, device="cpu")
print("\n示例3 - 指定类型和设备的全1张量:")
print(tensor3)
# 输出:
# tensor([[1, 1],
# [1, 1],
# [1, 1]])
# 示例4:创建需要梯度的全1张量(用于深度学习训练)
tensor4 = torch.ones(2, 2, requires_grad=True)
print("\n示例4 - 带梯度的全1张量:")
print(tensor4)
# 输出:
# tensor([[1., 1.],
# [1., 1.]], requires_grad=True)
size 参数的灵活写法
以下三种写法效果完全一致(均生成 [4,4] 全1张量),可根据场景选择:
torch.ones(4, 4) # 多个整数参数
torch.ones((4, 4)) # 元组参数(推荐,形状更清晰)
torch.ones([4, 4]) # 列表参数
与 torch.ones_like() 的区别
如果想生成“和某个已有张量形状相同”的全1张量,用 torch.ones_like() 更高效:
x = torch.randn(2, 5) # 已有张量,形状[2,5]
y = torch.ones_like(x) # 生成和x形状、类型、设备一致的全1张量
print(y.shape) # 输出:torch.Size([2,5])
在掩码场景中
torch.ones() 是构建掩码的“基础画布”:
- 先生成全1矩阵,再通过
torch.triu()调整哪些位置为1(需要屏蔽)、哪些为0(保留); - 对比:如果用
torch.zeros()做基础画布,后续需要反向调整(如下三角为1),不如torch.ones()直观。
总结
torch.ones(*size)核心作用是生成指定形状的全1张量,参数支持整数、元组/列表形式;- 可通过
dtype/device指定张量类型和存储设备,torch.ones_like()适配已有张量形状; - 在注意力掩码场景中,
torch.ones()是构建掩码矩阵的第一步,为后续triu/bool/masked_fill操作提供基础张量。