Running a conv autoencoder on MNIST dataset

Viewed 39

I have the following conv autoencoder implementation that I wrote:

class ConvolutionEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1    = nn.Conv2d(1,32,3,stride=1,padding=1,dilation=1)
        self.maxpool1 = nn.MaxPool2d(2,padding=0,return_indices=True)
        self.conv2    = nn.Conv2d(32,32,3,stride=1,padding=1,dilation=1)
        self.maxpool2 = nn.MaxPool2d(2,padding=0,return_indices=True)
        self.conv3    = nn.Conv2d(32,32,3,stride=1,padding=1,dilation=1)
        self.maxpool3 = nn.MaxPool2d(2,padding=0,return_indices=True)
        self.relu     = nn.ReLU()
        self.middle   = []
        self.maxpool1_index = None
        self.maxpool2_index = None
        self.maxpool3_index = None
    
    def forward(self, image):
      temp = self.conv1(image)
      temp = self.relu(temp)
      self.middle = [temp]
      temp, maxpool1_index = self.maxpool1(temp)
      self.middle.append(temp)
      temp = self.conv2(temp)
      temp = self.relu(temp)
      self.middle.append(temp)
      feature_map, maxpool2_index = self.maxpool2(temp)
      temp = self.conv3(temp)
      temp = self.relu(temp)
      self.middle.append(temp)
      feature_map, maxpool3_index = self.maxpool3(temp)
      self.maxpool1_index = maxpool1_index
      self.maxpool2_index = maxpool2_index
      self.maxpool3_index = maxpool3_index
      return feature_map

class ConvolutionDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.unmaxpool1    = nn.MaxUnpool2d(2,padding=0)
        self.conv1         = nn.ConvTranspose2d(32,32,3,stride=1,padding=1,dilation=1)
        self.unmaxpool2    = nn.MaxUnpool2d(2,padding=0)
        self.conv2         = nn.ConvTranspose2d(32,32,3,stride=1,padding=1,dilation=1)
        self.unmaxpool3    = nn.MaxUnpool2d(2,padding=0)
        self.conv3         = nn.ConvTranspose2d(32,1,3,stride=1,padding=1,dilation=1)
        self.relu          = nn.ReLU()
        self.sigmoid       = nn.Sigmoid()
        self.middle        = []

    def forward(self, feature_map, maxpool_index):
        temp = self.unmaxpool1(feature_map, maxpool_index[0])
        self.middle = [temp]
        temp = self.relu(temp)
        temp = self.conv1(temp)
        self.middle.append(temp)
        temp = self.unmaxpool2(temp, maxpool_index[1])
        self.middle.append(temp)
        temp = self.relu(temp)
        temp = self.conv2(temp)
        temp = self.unmaxpool3(temp, maxpool_index[2])
        self.middle.append(temp)
        temp = self.relu(temp)
        temp = self.conv3(temp)
        reconstructed_image = self.sigmoid(temp)
        return reconstructed_image

class ConvolutionAutoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = ConvolutionEncoder()
        self.decoder = ConvolutionDecoder()

    def forward(self, image):
        reconstructed_image = self.decoder(self.encoder(image), [self.encoder.maxpool3_index, self.encoder.maxpool2_index, self.encoder.maxpool1_index])
        return reconstructed_image

# Load train dataloder
train_data_transformed = torchvision.datasets.MNIST(root="/MNIST", train=True, download=True,transform=torchvision.transforms.ToTensor())
train_dataloader = DataLoader(train_data_transformed, batch_size=1024)

# Load test dataloder
test_data_transformed = torchvision.datasets.MNIST(root="/MNIST", train=False, download=True,transform=torchvision.transforms.ToTensor())
test_dataloader = DataLoader(test_data_transformed, batch_size=1024)

conv_autoencoder = ConvolutionAutoencoder()
conv_optimizer   = torch.optim.AdamW(conv_autoencoder.parameters())
conv_MSELoss     = nn.MSELoss()

for conv_epoch_idx in tqdm(range(epochs)):
    for conv_batch_idx, (imgs, _) in enumerate(train_dataloader):
        conv_optimizer.zero_grad()
        reconstructed = conv_autoencoder(imgs)

I'm getting the following error:

Expected shape of indices to be same as that of the input tensor ([1024, 32, 14, 14]) but got indices tensor with shape: [1024, 32, 7, 7]

As I understand I get it from temp = self.unmaxpool2(temp, maxpool_index[1]). I want it to have three Cov2d layers and three max pool layers. I'm trying to run this autoencoder on MNIST dataset. I guess the sizes that I provided do not match. I can't seem to figure out how can I fix the sizes so it will work. Is it possible to suggest how to solve this?

1 Answers

There is an error in the forward pass of your encoder. You named the output of the second maxpool feature_map instead of temp which is the name you use in the following layers.

Changing to this makes your code working as expected:

class ConvolutionEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1    = nn.Conv2d(1,32,3,stride=1,padding=1,dilation=1)
        self.maxpool1 = nn.MaxPool2d(2,padding=0,return_indices=True)
        self.conv2    = nn.Conv2d(32,32,3,stride=1,padding=1,dilation=1)
        self.maxpool2 = nn.MaxPool2d(2,padding=0,return_indices=True)
        self.conv3    = nn.Conv2d(32,32,3,stride=1,padding=1,dilation=1)
        self.maxpool3 = nn.MaxPool2d(2,padding=0,return_indices=True)
        self.relu     = nn.ReLU()
        self.middle   = []
        self.maxpool1_index = None
        self.maxpool2_index = None
        self.maxpool3_index = None
    
    def forward(self, image):
        temp = self.conv1(image)
        temp = self.relu(temp)
        self.middle = [temp]
        temp, maxpool1_index = self.maxpool1(temp)
        self.middle.append(temp)
        temp = self.conv2(temp)
        temp = self.relu(temp)
        self.middle.append(temp)
        temp, maxpool2_index = self.maxpool2(temp)
        temp = self.conv3(temp)
        temp = self.relu(temp)
        self.middle.append(temp)
        feature_map, maxpool3_index = self.maxpool3(temp)
        self.maxpool1_index = maxpool1_index
        self.maxpool2_index = maxpool2_index
        self.maxpool3_index = maxpool3_index
        return feature_map
Related