Pytorch Convolutional Autoencoders - pytorch

Question

How one construct decoder part of convolutional autoencoder? Suppose I have this
(input -> conv2d -> maxpool2d -> maxunpool2d -> convTranspose2d -> output):
# CIFAR images shape = 3 x 32 x 32
class ConvDAE(nn.Module):
def __init__(self):
super().__init__()
# input: batch x 3 x 32 x 32 -> output: batch x 16 x 16 x 16
self.encoder = nn.Sequential(
nn.Conv2d(3, 16, 3, stride=1, padding=1), # batch x 16 x 32 x 32
nn.ReLU(),
nn.BatchNorm2d(16),
nn.MaxPool2d(2, stride=2) # batch x 16 x 16 x 16
)
# input: batch x 16 x 16 x 16 -> output: batch x 3 x 32 x 32
self.decoder = nn.Sequential(
# this line does not work
# nn.MaxUnpool2d(2, stride=2, padding=0), # batch x 16 x 32 x 32
nn.ConvTranspose2d(16, 16, 3, stride=2, padding=1, output_padding=1), # batch x 16 x 32 x 32
nn.ReLU(),
nn.BatchNorm2d(16),
nn.ConvTranspose2d(16, 3, 3, stride=1, padding=1, output_padding=0), # batch x 3 x 32 x 32
nn.ReLU()
)
def forward(self, x):
print(x.size())
out = self.encoder(x)
print(out.size())
out = self.decoder(out)
print(out.size())
return out
Pytorch specific question: why can't I use MaxUnpool2d in decoder part. This gives me the following error:
TypeError: forward() missing 1 required positional argument: 'indices'
And the conceptual question: Shouldn't we do in decoder inverse of whatever we did in encoder? I saw some implementations and it seems they only care about the dimensions of input and output of decoder. Here and here are some examples.

Answer 1

For the torch part of the question, unpool modules have as a required positional argument the indices returned from the pooling modules which will be returned with return_indices=True. So you could do
class ConvDAE(nn.Module):
def __init__(self):
super().__init__()
# input: batch x 3 x 32 x 32 -> output: batch x 16 x 16 x 16
self.encoder = nn.Sequential(
nn.Conv2d(3, 16, 3, stride=1, padding=1), # batch x 16 x 32 x 32
nn.ReLU(),
nn.BatchNorm2d(16),
nn.MaxPool2d(2, stride=2, return_indices=True)
)
self.unpool = nn.MaxUnpool2d(2, stride=2, padding=0)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(16, 16, 3, stride=2, padding=1, output_padding=1),
nn.ReLU(),
nn.BatchNorm2d(16),
nn.ConvTranspose2d(16, 3, 3, stride=1, padding=1, output_padding=0),
nn.ReLU()
)
def forward(self, x):
print(x.size())
out, indices = self.encoder(x)
out = self.unpool(out, indices)
out = self.decoder(out)
print(out.size())
return out
As for the general part of the question, I don't think state of the art is to use a symmetric decoder part, as it has been shown that devonvolution/transposed convolution produces checkerboard effects and many approaches tend to use upsampling modules instead. You will find more info faster through PyTorch channels.