问题描述
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),5,5
nn.Conv2d(in_channels=32,out_channels=64,stride=2,# 64,3,3
nn.ReLU(True),stride=1),2,2
)
是否有像 nn.Sequential 这样的结构将模块并行放入其中？
我现在想定义类似的东西
self._mean_logvar_layers = nn.Parallel(
nn.Conv2d(in_channels=64,kernel_size=2,stride=1,padding=0),nn.Conv2d(in_channels=64,)
谁的输出应该是两个数据管道 - self._mean_logvar_layers 中的每个元素一个，然后可以提供给网络的其余部分。有点像多头网络。
我目前的实现：
self._mean_layer = nn.Conv2d(in_channels=64,padding=0)
self._logvar_layer = nn.Conv2d(in_channels=64,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,)
像这样使用它：
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,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
它只是将上一个模型的输出传递给下一个模型，就是这样。
,在@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,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,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),5,5
nn.Conv2d(in_channels=32,stride=2,# 64,3,3
nn.ReLU(True),stride=1),Split(
# notice out_channels are double of real desired out_channels
nn.Conv2d(in_channels=64,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 并在初始化时间定义不同的编码器。