Autoencoder model either oscillates or doesn't converge on MNIST dataset

Question

Already ran the code 3 months ago with intended results. Changed nothing. Tried troubleshooting by using codes from (several) earlier versions, including among the earliest (which definitely worked). The problem persists.

# 4 - Constructing the undercomplete architecture
class autoenc(nn.Module):
    def __init__(self, nodes = 100):
        super(autoenc, self).__init__() # inheritence
        self.full_connection0 = nn.Linear(784, nodes) # encoding weights
        self.full_connection1 = nn.Linear(nodes, 784) # decoding weights
        self.activation = nn.Sigmoid()

    def forward(self, x):
        x = self.activation(self.full_connection0(x)) # input encoding
        x = self.full_connection1(x) # output decoding
        return x



# 5 - Initializing autoencoder, squared L2 norm, and optimization algorithm
model = autoenc().cuda()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(),
                          lr = 1e-3, weight_decay = 1/2)



# 6 - Training the undercomplete autoencoder model
num_epochs = 500
batch_size = 32
length = int(len(trn_data) / batch_size)

loss_epoch1 = []

for epoch in range(num_epochs):
    train_loss = 0
    score = 0. 


    for num_data in range(length - 2):
        batch_ind = (batch_size * num_data)
        input = Variable(trn_data[batch_ind : batch_ind + batch_size]).cuda()

        # === forward propagation ===
        output = model(input)
        loss = criterion(output, trn_data[batch_ind : batch_ind + batch_size])

        # === backward propagation ===
        loss.backward()

        # === calculating epoch loss ===
        train_loss += np.sqrt(loss.item())
        score += 1. #<- add for average loss error instead of total
        optimizer.step()

    loss_calculated = train_loss/score
    print('epoch: ' + str(epoch + 1) + '   loss: ' + str(loss_calculated))
    loss_epoch1.append(loss_calculated)

When plotting the loss now, it oscillates oscillates wildly (at lr = 1e-3). Whereas 3 months ago, it was steadily converging (at lr = 1e-3).

Can't upload pictures yet due to recently created account.

How it looks like now.

Though this is when I reduce the learning rate to 1e-5. When it's at 1e-3, it's just all over the places.

How it should look like, and used to look like at lr = 1e-3.

Answer 1

You should do optimizer.zero_grad() before you do loss.backward() since the gradients accumulate. This is most likely causing the issue.

The general order to be followed during training phase :

optimizer.zero_grad()
output = model(input)
loss = criterion(output, label)
loss.backward()
optimizer.step()

Also, the value of weight decay used (1 / 2) was causing an issue.