通过顺序容器展平PyTorch构建层

Ste*_*ing 8 python conv-neural-network pytorch

我试图通过PyTorch的顺序容器构建一个cnn,我的问题是我无法弄清楚如何展平图层.

main = nn.Sequential()
self._conv_block(main, 'conv_0', 3, 6, 5)
main.add_module('max_pool_0_2_2', nn.MaxPool2d(2,2))
self._conv_block(main, 'conv_1', 6, 16, 3)
main.add_module('max_pool_1_2_2', nn.MaxPool2d(2,2)) 
main.add_module('flatten', make_it_flatten)
Run Code Online (Sandbox Code Playgroud)

我应该在"make_it_flatten"中加入什么?我试图压扁主但是它不起作用,主要不存在调用视图的东西

main = main.view(-1, 16*3*3)
Run Code Online (Sandbox Code Playgroud)

cle*_*ros 17

这可能不是您正在寻找的,但您可以简单地创建自己的nn.Module扁平化任何输入,然后您可以将其添加到nn.Sequential()对象:

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size()[0], -1)
Run Code Online (Sandbox Code Playgroud)

x.size()[0]将选择批次暗淡,并且-1将计算所有剩余的变暗,以适应元件的数量,从而压扁任何张量/变量.

并在nn.Sequential以下方面使用它:

main = nn.Sequential()
self._conv_block(main, 'conv_0', 3, 6, 5)
main.add_module('max_pool_0_2_2', nn.MaxPool2d(2,2))
self._conv_block(main, 'conv_1', 6, 16, 3)
main.add_module('max_pool_1_2_2', nn.MaxPool2d(2,2)) 
main.add_module('flatten', Flatten())
Run Code Online (Sandbox Code Playgroud)