I am trying to build a simple classification network by PyTorch, but I do not how to flatten the convolution layer.
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
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))
#how could I flatten the convolution layer?
self._main = main
def forward(self, x):
return x
def _conv_block(main, name, inp_filter_size, out_filter_size, kernal_size):
main.add_module('{}-{}.{}.conv'.format(name, inp_filter_size, out_filter_size),
nn.Conv2d(inp_filter_size, out_filter_size, kernal_size, 1, 1))
main.add_module('{}-{}.batchnorm'.format(name, out_filter_size), nn.BatchNorm2d(out_filter_size))
main.add_module('{}-{}.relu'.format(name, out_filter_size), nn.ReLU())
Thanks