Thank you for suggestion! I did try BatchNorm after every layer and while it got slightly better it still doesn’t match the results of the shallower network. Here are the new convergence curves:
LossPlot_BatchNorm2880×1630 87.6 KB
You can see that the new BatchNorm curve (model 3) is better than model 2 but still not as good as the shallower model #1. I am not sure why! Here’s my code, just to confirm that I am doing the BatchNorm correctly:
# model 3 - BatchNorm
class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        # encoder
        self.enc0 = nn.Linear(in_features=512, out_features=512)
        self.e0 = nn.BatchNorm1d(392, affine=True)
        self.enc1 = nn.Linear(in_features=512, out_features=256)
        self.e1 = nn.BatchNorm1d(256, affine=True)
        self.enc2 = nn.Linear(in_features=256, out_features=256)
        self.e2 = nn.BatchNorm1d(256, affine=True)
        self.enc3 = nn.Linear(in_features=256, out_features=128)
        self.e3 = nn.BatchNorm1d(128, affine=True)
        self.enc4 = nn.Linear(in_features=128, out_features=128)
        self.e4 = nn.BatchNorm1d(128, affine=True)
        self.enc5 = nn.Linear(in_features=128, out_features=64)
        self.e5 = nn.BatchNorm1d(64, affine=True)
        self.enc6 = nn.Linear(in_features=64, out_features=64)
        self.e6 = nn.BatchNorm1d(64, affine=True)
 
        # decoder 
        self.dec0 = nn.Linear(in_features=64, out_features=64)
        self.d0 = nn.BatchNorm1d(64, affine=True)
        self.dec1 = nn.Linear(in_features=64, out_features=128)
        self.d1 = nn.BatchNorm1d(128, affine=True)
        self.dec2 = nn.Linear(in_features=128, out_features=128)
        self.d2 = nn.BatchNorm1d(128, affine=True)
        self.dec3 = nn.Linear(in_features=128, out_features=256)
        self.d3 = nn.BatchNorm1d(256, affine=True)
        self.dec4 = nn.Linear(in_features=256, out_features=256)
        self.d4 = nn.BatchNorm1d(256, affine=True)
        self.dec5 = nn.Linear(in_features=256, out_features=512)
        self.d5 = nn.BatchNorm1d(392, affine=True)
        self.dec6 = nn.Linear(in_features=512, out_features=512)

    def forward(self, x):
        # encoder
        x = F.relu(self.enc0(x))
        x = self.e0(x)
        x = F.relu(self.enc1(x))
        x = self.e1(x)
        x = F.relu(self.enc2(x))
        x = self.e2(x)
        x = F.relu(self.enc3(x))
        x = self.e3(x)
        x = F.relu(self.enc4(x))
        x = self.e4(x)
        x = F.relu(self.enc5(x))
        x = self.e5(x)
        x = F.relu(self.enc6(x))
        x = self.e6(x)
        
        # decoder
        x = F.relu(self.dec0(x))
        x = self.d0(x)
        x = F.relu(self.dec1(x))
        x = self.d1(x)
        x = F.relu(self.dec2(x))
        x = self.d2(x)
        x = F.relu(self.dec3(x))
        x = self.d3(x)
        x = F.relu(self.dec4(x))
        x = self.d4(x)
        x = F.relu(self.dec5(x))
        x = self.d5(x)
        x = self.dec6(x) # no RELU on the last one

        return x

It’s very strange. I thought it would work better! What am I doing wrong?