00200 PyTorch - 常用函数 ubuntu | LuYF-Lemon-love の Blog

前言

简单的介绍PyTorch - 常用函数。

Operating System: Ubuntu 22.04.4 LTS

独立函数

torch.unsqueeze

api link: https://pytorch.org/docs/stable/generated/torch.unsqueeze.html#torch.unsqueeze

unsqueeze 是 PyTorch 中的一个函数，用于增加一个维度到给定张量（tensor）的指定位置。这个函数常用于在神经网络中处理数据的维度，确保数据能够正确地传递到网络的下一层。

以下是 unsqueeze 函数的基本用法：

1	torch.unsqueeze(input, dim)

input：要处理的张量。
dim：要在其中插入新维度的位置。

如果 dim 是一个整数，那么新维度将会被插入到 input 张量的指定位置。注意，PyTorch 中的维度索引是从0开始的。如果 dim 是负数，则新维度将会被添加到从末尾开始的指定位置。

这里有一个例子：

import torch
x = torch.tensor([1, 2, 3, 4])  # 创建一个一维张量
print(x.size())  # 输出: torch.Size([4])
y = x.unsqueeze(0)  # 在第0维（最外层）增加一个维度
print(y.size())  # 输出: torch.Size([1, 4])
z = x.unsqueeze(1)  # 在第1维增加一个维度
print(z.size())  # 输出: torch.Size([4, 1])

在上面的例子中，x 是一个一维张量。通过 unsqueeze(0)，我们在最外层增加了一个维度，使得 y 成为了一个 2D 张量，形状为 [1, 4]。类似地，通过 unsqueeze(1)，我们在第二个位置增加了一个维度，使得 z 成为了一个形状为 [4, 1] 的 2D 张量。

unsqueeze 常用于将一个一维张量转换为二维张量（例如，批量大小为1的情况），或者将一个二维张量转换为三维张量，其中增加的维度通常用于表示批次大小、通道数等。

torch.repeat_interleave

api link: https://pytorch.org/docs/stable/generated/torch.repeat_interleave.html#torch.repeat_interleave

torch.repeat_interleave 是PyTorch中用于张量操作的函数，它可以沿着指定的维度重复张量中的元素。这个函数对于创建具有特定模式的重复数据非常有用。

以下是 torch.repeat_interleave 的基本用法：

1	torch.repeat_interleave(input, repeats, dim=None)

input: 要重复的输入张量。
repeats: 一个1-D张量，指定了 input 张量在 dim 维度上每个元素的重复次数。
dim: 要重复的维度，如果 dim 是None，那么 input 会被展平成一维张量后再进行重复。

下面是一个简单的例子：

import torch
x = torch.tensor([1, 2, 3])
repeats = torch.tensor([2, 3, 4])
result = torch.repeat_interleave(x, repeats, dim=0)
print(result)

这个例子中，张量 x 的元素 [1, 2, 3] 将分别被重复 2 次、3 次和 4 次，输出结果将是：

1	tensor([1, 1, 2, 2, 2, 3, 3, 3, 3])

如果 dim 参数未指定，input 张量首先会被展平，然后根据 repeats 张量中的值进行重复。

这个函数在进行数据预处理和创建批量数据时非常有用，尤其是在需要根据不同的样本长度进行重复操作时。

Module 函数

register_buffer

api link: https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer

在PyTorch中，register_buffer是一个在模型中注册一个缓冲区（buffer）的方法，这个缓冲区在模型保存和加载时会被包含进去，但它不是一个模型参数，也就是说不会被优化器更新。

缓冲区通常用于存储一些不经常变化的数据，比如批量归一化的均值和方差，或者是在训练RNN时保存隐藏状态等。

下面是register_buffer的一些基本用法：

定义

1	def register_buffer(self, name, tensor, persistent=True):

name: 缓冲区的名称，通过这个名称可以在模型中访问这个缓冲区。
tensor: 要注册的Tensor。
persistent: 当设置为True时，这个缓冲区在调用state_dict()时会被包含，并且在模型保存和加载时会被持久化。

使用示例

import torch
import torch.nn as nn
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.register_buffer('running_mean', torch.zeros(100))
    def forward(self, x):
        # 使用running_mean进行一些操作
        pass
model = MyModel()
print(model.state_dict())