【Torch API】pytorch 中torch.ones_like和torch.zeros_like函数详解

devil_son1234

已于 2023-03-17 15:48:21 修改

阅读量829

点赞数

文章标签： python

于 2023-03-17 15:48:13 首次发布

原文链接：https://blog.youkuaiyun.com/tszupup/article/details/108130721

版权

本文介绍了torch.ones_like和torch.zeros_like函数，这两个函数用于创建与给定张量形状相同但元素全为1或0的新张量。用户可以自定义数据类型、布局、设备及是否计算梯度，以适应不同的计算需求。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch.ones_like函数和torch.zeros_like函数的基本功能是根据给定张量，生成与其形状相同的全1张量或全0张量，示例如下：

input = torch.rand(2, 3)
print(input)
# 生成与input形状相同、元素全为1的张量
a = torch.ones_like(input)
print(a)
# 生成与input形状相同、元素全为0的张量
b = torch.zeros_like(input)
print(b)

效果如下：

tensor([[0.0881, 0.9002, 0.7084],
        [0.3313, 0.2736, 0.0894]])
tensor([[1., 1., 1.],
        [1., 1., 1.]])
tensor([[0., 0., 0.],
        [0., 0., 0.]])

我们进一步看一下这两个函数在源码中是怎样定义的。

torch.ones_like函数：

@overload
def ones_like(self: Tensor, *, dtype: _dtype=None, layout: layout=strided, device: Union[_device, str, None]=None, requires_grad:_bool=False) -> Tensor: ...

torch.zeros_like函数：

@overload
def zeros_like(self: Tensor, *, dtype: _dtype=None, layout: layout=strided, device: Union[_device, str, None]=None, requires_grad:_bool=False) -> Tensor: ...

可以看到，在这两个函数中，我们还可以指定数据类型、设备、是否计算梯度等信息，可以结合具体场景灵活使用。