Pytorch 两个函数 new_ones 和 ones

在本文中，我们将介绍 PyTorch 中的两个函数 new_ones 和 ones，它们都可以用来创建指定形状的张量，并填充为指定的值。尽管它们的功能相似，但它们的用法和效果有所不同。

ones 函数

ones 函数是 PyTorch 提供的一个常用函数，它可以创建一个指定形状的张量，并将所有元素初始化为 1。ones 函数的用法如下：

import torch

shape = (2, 3)
tensor_ones = torch.ones(shape)
print(tensor_ones)

运行上述代码，将输出一个形状为 (2, 3) 的张量，并且所有元素的值都为 1。输出结果如下：

tensor([[1., 1., 1.],
        [1., 1., 1.]])

我们可以看到，使用 ones 函数创建的张量中的所有元素都被设置为 1。这在一些需要初始化为特定值的情况下非常有用，例如权重初始化等。

new_ones 函数

new_ones 函数是 PyTorch 提供的另一个函数，它也可以创建一个指定形状的张量，并将所有元素初始化为 1。但与 ones 不同的是，new_ones 函数可以在指定的设备上为张量分配内存。new_ones 函数的用法如下：

import torch

shape = (2, 3)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tensor_new_ones = torch.ones(shape, device=device)
print(tensor_new_ones)

上述代码中，我们首先根据当前系统的可用设备选择将张量分配到 GPU 还是 CPU。然后，我们使用 new_ones 创建一个形状为 (2, 3) 的张量，并将该张量分配到选定的设备上。输出结果如下：

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')

我们可以看到，使用 new_ones 创建的张量与 ones 函数相同，但是该张量被分配到了 CUDA 设备上（如果可用）。这对于在 GPU 上进行计算时非常有用。

示例说明

为了更好地理解 ones 和 new_ones 的使用，我们以一个简单的示例来说明它们的区别。假设我们有一个形状为 (3, 3) 的张量，并使用 ones 和 new_ones 分别创建两个相同形状的张量。

import torch

shape = (3, 3)
tensor_ones = torch.ones(shape)
tensor_new_ones = torch.ones(shape, device="cuda" if torch.cuda.is_available() else "cpu")

在上述示例中，我们使用 ones 创建了一个形状为 (3, 3) 的张量 tensor_ones，并使用 new_ones 创建了一个相同形状的张量 tensor_new_ones，并分配到了合适的设备上。

接下来，我们可以比较这两个张量：

print("Tensor created using ones:")
print(tensor_ones)
print("Tensor created using new_ones:")
print(tensor_new_ones)

输出结果如下：

Tensor created using ones:
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])
Tensor created using new_ones:
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')

可以看到，两个张量的输出结果相同，都是一个形状为 (3, 3) 的张量，并且所有元素的值都为 1。但 tensor_new_ones 还显示了它被分配到了 CUDA 设备上。