是否存在与 torch.nn.Sequential 等效的“Split”？

Posted 2023-03-27

【中文标题】是否存在与 torch.nn.Sequential 等效的“Split”？

【英文标题】：Is there a `Split` equivalent to torch.nn.Sequential?

【发布时间】：2021-04-26 02:10:23

【问题描述】：

Sequential 块的示例代码是

self._encoder = nn.Sequential(
        # 1, 28, 28
        nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=3, padding=1),
        # 32, 10, 10 = 16, (1//3)(28 + 2 * 1 - 3) + 1, (1//3)(28 + 2*1 - 3) + 1
        nn.ReLU(True),
        nn.MaxPool2d(kernel_size=2, stride=2),
        # 32, 5, 5
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
        # 64, 3, 3
        nn.ReLU(True),
        nn.MaxPool2d(kernel_size=2, stride=1),
        # 64, 2, 2
)

---

是否有像nn.Sequential 这样的结构将模块并行放入其中？

我现在想定义类似的东西

self._mean_logvar_layers = nn.Parallel(
    nn.Conv2d(in_channels=64, out_channels=64, kernel_size=2, stride=1, padding=0),
    nn.Conv2d(in_channels=64, out_channels=64, kernel_size=2, stride=1, padding=0),
)

其输出应该是两个数据管道 - self._mean_logvar_layers 中的每个元素一个管道，然后可馈送到网络的其余部分。有点像多头网络。

---

我目前的实现：

self._mean_layer = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=2, stride=1, padding=0)
self._logvar_layer = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=2, stride=1, padding=0)

和

def _encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    for i, layer in enumerate(self._encoder):
        x = layer(x)

    mean_output = self._mean_layer(x)
    logvar_output = self._logvar_layer(x)

    return mean_output, logvar_output

我想将并行构造视为一个层。

这在 PyTorch 中可行吗？

【问题讨论】：

并行是并行运行（同时）还是只输出两个值但按顺序工作？如果是第一种情况，这种并行化究竟意味着什么（因为您可以通过一些.to 调用和不同的设备轻松完成，请参阅here）。

@SzymonMaszke 我的意思是不要考虑后端技术问题。我想要类似于我当前实现的行为，但是将它包装在一个允许我在网络架构中定义“拆分”的构造中。最好并行运行会很好，但我对 GPU 并行最佳实践的了解要少得多，因此希望让 pytorch 为我做魔法。

只需使用一层，后跟torch.split。请参阅此处：pytorch.org/docs/stable/generated/torch.split.html 在这种情况下，您将使用具有 128 个输出通道的 conv，并沿轴 1 分成两个大小为 64 的部分（假设为 channels_first 数据格式）。这可能不是一个通用解决方案，所以我不愿将其作为答案发布，但它适用于大多数常见情况，如 VAE。

感谢添加。我也可以尝试写一个答案，但我对 pytorch 不是很熟悉——我使用的是 tensorflow，这就是我要做的。 tf.split 很好，因为它还允许您简单地指定拆分的 number （在本例中为 2）并计算输出大小。看起来火炬不支持这个。至于 Gulzars 最后的评论：实际上，一个大小为 128 的 FC 层在数学上与接收相同输入的两个大小为 64 的层的串联完全相同（卷积类似）。这就是为什么这样做“没问题”。

@Gulzar 你有两次N 输入和M 输出（均值和方差）。每个M 参见N。在N->2M 的情况下，每个M 也会看到N 输入。和@xdurch0 说的完全一样，Ms 之间也没有联系。对于 FC 来说，它就像多类逻辑回归（没有激活），每个权重向量都独立于其他权重向量（与所指出的卷积相同）。

【参考方案1】：

顺序拆分

你可以做的是创建一个Parallel 模块（虽然我会以不同的方式命名它，因为这意味着这段代码实际上是并行运行的，Split 可能是一个好名字），如下所示：

class Parallel(torch.nn.Module):
    def __init__(self, *modules: torch.nn.Module):
        super().__init__()
        self.modules = modules

    def forward(self, inputs):
        return [module(inputs) for module in self.modules]

现在您可以根据需要定义它：

self._mean_logvar_layers = Parallel(
    nn.Conv2d(in_channels=64, out_channels=64, kernel_size=2, stride=1, padding=0),
    nn.Conv2d(in_channels=64, out_channels=64, kernel_size=2, stride=1, padding=0),
)

并像这样使用它：

mean, logvar = self._mean_logvar_layers(x)

一层并拆分

正如@xdurch0 所建议的那样，我们可以使用单层并跨通道拆分，使用此模块：

class Split(torch.nn.Module):
    def __init__(self, module, parts: int, dim=1):
        super().__init__()
        self.parts
        self.dim = dim
        self.module = module

    def forward(self, inputs):
        output = self.module(inputs)
        chunk_size = output.shape[self.dim] // self.parts
        return torch.split(output, chunk_size, dim=self.dim)

这在你的神经网络中（注意128 通道，这些通道将被分成2 部分，每个部分的大小为64）：

self._mean_logvar_layers = Split(
    nn.Conv2d(in_channels=64, out_channels=128, kernel_size=2, stride=1, padding=0),
    parts=2,
)

并像以前一样使用它：

mean, logvar = self._mean_logvar_layers(x)

为什么采用这种方法？

一切都将一举完成，而不是按顺序完成，因此速度更快，但如果您没有足够的 GPU 内存，可能会太宽。

它可以与 Sequential 一起使用吗？

是的，它仍然是一个层。但下一层必须使用tuple(torch.Tensor, torch.Tensor) 作为输入。

Sequential也是一层，挺简单的，看forward：

def forward(self, inp):
    for module in self:
        inp = module(inp)
    return inp

它只是将前一个模型的输出传递给下一个模型，就是这样。

【讨论】：

谢谢！ Split 模块可以像一个层一样被输入到nn.Sequential 中吗？

@Gulzar 查看我的编辑，但在这种情况下为什么需要它？

所有这些实际上是为了能够简化我的 VAE 以便能够在编码器中包含 mean 和 logvar。我希望它能够以这种方式覆盖每个数据集的 _encoder 和 _decoder，但仍使用相同的 VAE 包装器代码。我正在 MNIST 上测试我的基础架构，但实际上使用的是不适合调试的独特数据。

@Gulzar 您可以创建许多不同的层（例如将tuple 作为输入并在每个项目上应用不同的modules，类似于上面答案中的第一个Parallel）。总而言之，它是 Python，所以list comprehensions 和相关的东西都可以工作，你可以按照你想要的任何方式来构建它（尽管它是否值得抽象）。

让我们continue this discussion in chat.

【参考方案2】：

根据@Szymon Maszke 的精彩回答，这里是完整的相关代码，经过所有扩充：

class Split(torch.nn.Module):
    """
    https://***.com/questions/65831101/is-there-a-parallel-equivalent-to-toech-nn-sequencial#65831101

    models a split in the network. works with convolutional models (not FC).
    specify out channels for the model to divide by n_parts.
    """
    def __init__(self, module, n_parts: int, dim=1):
        super().__init__()
        self._n_parts = n_parts
        self._dim = dim
        self._module = module

    def forward(self, inputs):
        output = self._module(inputs)
        chunk_size = output.shape[self._dim] // self._n_parts
        return torch.split(output, chunk_size, dim=self._dim)


class Unite(torch.nn.Module):
    """
    put this between two Splits to allow them to coexist in sequence.
    """
    def __init__(self):
        super(Unite, self).__init__()

    def forward(self, inputs):
        return torch.cat(inputs, dim=1)

及用法：

class VAEConv(VAEBase):
    ...
    ...
    ...

    def __init__():
        ...
        ...
        ...
        self._encoder = nn.Sequential(
        # 1, 28, 28
        nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=3, padding=1),
        # 32, 10, 10 = 16, (1//3)(28 + 2 * 1 - 3) + 1, (1//3)(28 + 2*1 - 3) + 1
        nn.ReLU(True),
        nn.MaxPool2d(kernel_size=2, stride=2),
        # 32, 5, 5
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
        # 64, 3, 3
        nn.ReLU(True),
        nn.MaxPool2d(kernel_size=2, stride=1),
        Split(
                # notice out_channels are double of real desired out_channels
                nn.Conv2d(in_channels=64, out_channels=128, kernel_size=2, stride=1, padding=0),
                n_parts=2,
        ),
        Unite(),
        Split(
                nn.Flatten(start_dim=1, end_dim=-1),
                n_parts=2
        ),
    )

    def _encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        for i, layer in enumerate(self._encoder):
            x = layer(x)

        mean_output, logvar_output = x
        return mean_output, logvar_output

这现在允许子类化 VAE 并在 init 时间定义不同的编码器。