一个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中这是可行的吗?
回答:
Sequential的分裂
你可以创建一个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)
单层和分裂
如@***所建议,我们可以使用单一层并按通道进行分裂,使用以下模块:
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
它只是将前一个模型的输出传递给下一个模型,仅此而已。