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?