[Python] CNN Autoencoder takes a very long time to train

I have been training a CNN Autoencoder on binary images (pixels are either 0 or 1) of size 64x64. The model is shown below:

import torch
import torch.nn as nn
import torch.nn.functional as F

class CNNAutoencoder(nn.Module):
"""
Input shape: (B, 1, 64, 64)
"""
def __init__(self, grid_size=64):
super().__init__()
self.grid_size = grid_size

# 1) Encoder layers
self.enc_conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
self.bn_enc1 = nn.BatchNorm2d(16)
self.enc_conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
self.bn_enc2 = nn.BatchNorm2d(32)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

# After two pools => shape is (32, grid_size//4, grid_size//4)
flat_dim = 32 * (grid_size // 4) * (grid_size // 4)
latent_dim = 128

# For linear layers, we can use BatchNorm1d:
self.fc_enc = nn.Linear(flat_dim, latent_dim)
self.bn_fc_enc = nn.BatchNorm1d(latent_dim)

# 2) Decoder layers
self.fc_dec = nn.Linear(latent_dim, flat_dim)
self.bn_fc_dec = nn.BatchNorm1d(flat_dim)

self.dec_tconv1 = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
self.bn_dec1 = nn.BatchNorm2d(16)
self.dec_tconv2 = nn.ConvTranspose2d(16, 1, kernel_size=2, stride=2)

def encoder(self, x):
# x => (B,1,64,64)
x = self.enc_conv1(x) # => (B,16,64,64)
x = self.bn_enc1(x)
x = F.relu(x)
x = self.pool(x) # => (B,16,32,32)

x = self.enc_conv2(x) # => (B,32,32,32)
x = self.bn_enc2(x)
x = F.relu(x)
x = self.pool(x) # => (B,32,16,16)

# Flatten
x = x.view(x.size(0), -1) # => (B, flat_dim=32*16*16)
x = self.fc_enc(x) # => (B, latent_dim=128)
x = self.bn_fc_enc(x)
x = F.relu(x)
return x

def decoder(self, z):
# z => (B,128)
x = self.fc_dec(z) # => (B, flat_dim)
x = self.bn_fc_dec(x)
x = F.relu(x)

# Reshape to (B, 32, 16, 16)
B = x.size(0)
x = x.view(B, 32, self.grid_size // 4, self.grid_size // 4)

x = self.dec_tconv1(x) # => (B,16,32,32)
x = self.bn_dec1(x)
x = F.relu(x)

# Final upsample => (B,1,64,64)
x = self.dec_tconv2(x)
x = torch.sigmoid(x)
return x

def forward(self, x):
z = self.encoder(x)
return self.decoder(z)


Train set includes 1079156 samples. I'm using batch size 128 and 7 epochs. Right now, 1 epoch takes approximately 3 hours to train. Does anyone have an idea what might be the problem?

I tried to switch between BatchNorm and LayerNorm/GroupNorm but the performance didn't change. I also tried early stopping, but this doesn't solve the problem of timing per epoch. For your reference, I've pasted my training code below:

train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False)

# Model
model = CNNAutoencoder(grid_size=64).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.BCELoss()

best_val_loss = float('inf')
model_fname = f"best_{rt}_fold{fold}.pth"

global_step = 0

for ep in range(num_epochs):
if global_step >= max_steps:
print(f"Reached {max_steps} total steps; stopping early.")
break

t0 = time.time()

#############################
# Train Loop (per epoch)
#############################
model.train()
total_loss = 0.0
for x_in, x_tgt in train_loader:
x_in = x_in.to(device)
x_tgt = x_tgt.to(device)

optimizer.zero_grad()
out = model(x_in)
loss = loss_fn(out, x_tgt)
loss.backward()
optimizer.step()

total_loss += loss.item() * x_in.size(0)

global_step += 1 # increment step count
if global_step >= max_steps:
print(f"Reached {max_steps} total steps; stopping in mid-epoch.")
break

train_epoch_loss = total_loss / len(train_loader.dataset)

# If we already hit max_steps, break out
if global_step >= max_steps:
break

#############################
# Validation Loop
#############################
model.eval()
val_loss_sum = 0.0
with torch.no_grad():
for x_in, x_tgt in val_loader:
x_in, x_tgt = x_in.to(device), x_tgt.to(device)
out = model(x_in)
loss = loss_fn(out, x_tgt)
val_loss_sum += loss.item() * x_in.size(0)

current_val_loss = val_loss_sum / len(val_loader.dataset)
dt = time.time() - t0

print(f"Epoch [{ep+1}/{num_epochs}] => "
f"train_loss={train_epoch_loss:.4f}, val_loss={current_val_loss:.4f}, dt={dt:.2f}s, "
f"global_step={global_step}")

# Save best model
if current_val_loss < best_val_loss:
best_val_loss = current_val_loss
torch.save(model.state_dict(), model_fname)```

Continue reading...