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 中可行吗?
你可以做的是创建一个
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 内存,可能会太宽。
是的,它仍然是一层。但下一层必须使用
tuple(torch.Tensor, torch.Tensor)
作为输入。
Sequential
也是一层,很简单的,我们看看forward
:
def forward(self, inp):
for module in self:
inp = module(inp)
return inp
它只是将前一个模型的输出传递到下一个模型,仅此而已。
遵循 @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 时间定义不同的编码器。