pytorch 是否有一个'分裂'相当于torch.nn.顺序?

wnrlj8wa  于 2023-10-20  发布在  其他
关注(0)|答案(2)|浏览(136)

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中可行吗?

qf9go6mv

qf9go6mv1#

顺序拆分

你可以做的是创建一个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

它只是将输出从前一个模型传递到下一个模型,仅此而已。

thtygnil

thtygnil2#

在@Szymon Maszke的伟大回答之后,这里是完整的相关代码,经过所有的增强:

class Split(torch.nn.Module):
    """
    https://stackoverflow.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时定义不同的编码器。

相关问题