machine-learningdeep-learningpytorchvgg-net

The training loss of vgg16 implemented in pytorch does not decrease


I want to try some toy examples in pytorch, but the training loss does not decrease in the training.

Some info is provided here:

The code is as follows

# encoding: utf-8

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision

import numpy as np


class VGG16(torch.nn.Module):
    def __init__(self, n_classes):
        super(VGG16, self).__init__()

        # construct model
        self.conv1_1 = nn.Conv2d(3, 64, 3, padding=1)
        self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1)
        self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1)
        self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1)
        self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1)
        self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1)
        self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1)

        self.fc6 = nn.Linear(512, 512)
        self.fc7 = nn.Linear(512, 512)
        self.fc8 = nn.Linear(512, n_classes)

    def forward(self, x):
        x = F.relu(self.conv1_1(x))
        x = F.relu(self.conv1_2(x))
        x = F.max_pool2d(x, (2, 2))

        x = F.relu(self.conv2_1(x))
        x = F.relu(self.conv2_2(x))
        x = F.max_pool2d(x, (2, 2))

        x = F.relu(self.conv3_1(x))
        x = F.relu(self.conv3_2(x))
        x = F.relu(self.conv3_3(x))
        x = F.max_pool2d(x, (2, 2))

        x = F.relu(self.conv4_1(x))
        x = F.relu(self.conv4_2(x))
        x = F.relu(self.conv4_3(x))
        x = F.max_pool2d(x, (2, 2))

        x = F.relu(self.conv5_1(x))
        x = F.relu(self.conv5_2(x))
        x = F.relu(self.conv5_3(x))
        x = F.max_pool2d(x, (2, 2))

        x = x.view(-1, self.num_flat_features(x))

        x = F.relu(self.fc6(x))
        x = F.relu(self.fc7(x))
        x = self.fc8(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]

        num_features = 1
        for s in size:
            num_features *= s
        return num_features


if __name__ == '__main__':

    BATCH_SIZE = 128
    LOG_INTERVAL = 5

    # data
    transform = transforms.Compose([
        transforms.ToTensor()
    ])

    trainset = torchvision.datasets.CIFAR100(
        root='./data',
        train=True,
        download=True,
        transform=transform
    )

    testset = torchvision.datasets.CIFAR100(
        root='./data',
        train=False,
        download=True,
        transform=transform
    )

    trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True)
    testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE, shuffle=False)

    # model
    vgg16 = VGG16(100)
    vgg16.cuda()

    # optimizer
    optimizer = optim.SGD(vgg16.parameters(), lr=0.01)

    # loss
    criterion = nn.CrossEntropyLoss()

    print('———— Train Start —————')
    for epoch in range(20):
        running_loss = 0.
        for step, (batch_x, batch_y) in enumerate(trainloader):
            batch_x, batch_y = batch_x.cuda(), batch_y.cuda()

            # 
            optimizer.zero_grad()

            output = vgg16(batch_x)
            loss = criterion(output, batch_y)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            if step % LOG_INTERVAL == 0:
                print('[%d, %4d] loss: %.4f' % (epoch, step, running_loss / LOG_INTERVAL))
                running_loss = 0.

        def test():
            print('———— Test Start ————')
            correct = 0
            total = 0

            # 
            with torch.no_grad():
                for test_x, test_y in testloader:
                    images, labels = test_x.cuda(), test_y.cuda()
                    output = vgg16(images)
                    _, predicted = torch.max(output.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()

            accuracy = 100 * correct / total
            print('Accuracy of the network is: %.4f %%' % accuracy)
            print('———— Test Finish ————')

        test()

    print('———— Train Finish —————')

The loss stays around 4.6060 and never decrease. I have tried different learning rate but does not work.


Solution

  • I have noticed that you are not using Batch normalization in between your convolution layers. I have added batch normalization layers and it seems to work. Following is the modified code:

    class VGG16(torch.nn.Module):
    def __init__(self, n_classes):
        super(VGG16, self).__init__()
    
        # construct model
        self.conv1_1 = nn.Conv2d(3, 64, 3, padding=1)
        self.conv11_bn = nn.BatchNorm2d(64)
        self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv12_bn = nn.BatchNorm2d(64)
        self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv21_bn = nn.BatchNorm2d(128)
        self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1)
        self.conv22_bn = nn.BatchNorm2d(128)
        self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1)
        self.conv31_bn = nn.BatchNorm2d(256)
        self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1)
        self.conv32_bn = nn.BatchNorm2d(256)
        self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1)
        self.conv33_bn = nn.BatchNorm2d(256)
        self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1)
        self.conv41_bn = nn.BatchNorm2d(512)
        self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv42_bn = nn.BatchNorm2d(512)
        self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv43_bn = nn.BatchNorm2d(512)
        self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv51_bn = nn.BatchNorm2d(512)        
        self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv52_bn = nn.BatchNorm2d(512)
        self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1)
        self.conv53_bn = nn.BatchNorm2d(512)
    
        self.fc6 = nn.Linear(512, 512)
        self.fc7 = nn.Linear(512, 512)
        self.fc8 = nn.Linear(512, n_classes)
    
    def forward(self, x):
        x = F.relu(self.conv11_bn(self.conv1_1(x)))
        x = F.relu(self.conv12_bn(self.conv1_2(x)))
        x = F.max_pool2d(x, (2, 2))
    
        x = F.relu(self.conv22_bn(self.conv2_1(x)))
        x = F.relu(self.conv21_bn(self.conv2_2(x)))
        x = F.max_pool2d(x, (2, 2))
    
        x = F.relu(self.conv31_bn(self.conv3_1(x)))
        x = F.relu(self.conv32_bn(self.conv3_2(x)))
        x = F.relu(self.conv33_bn(self.conv3_3(x)))
        x = F.max_pool2d(x, (2, 2))
    
        x = F.relu(self.conv41_bn(self.conv4_1(x)))
        x = F.relu(self.conv42_bn(self.conv4_2(x)))
        x = F.relu(self.conv43_bn(self.conv4_3(x)))
        x = F.max_pool2d(x, (2, 2))
    
        x = F.relu(self.conv51_bn(self.conv5_1(x)))
        x = F.relu(self.conv52_bn(self.conv5_2(x)))
        x = F.relu(self.conv53_bn(self.conv5_3(x)))
        x = F.max_pool2d(x, (2, 2))
    
        x = x.view(-1, self.num_flat_features(x))
    
        x = F.relu(self.fc6(x))
        x = F.relu(self.fc7(x))
        x = self.fc8(x)
        return x
    

    However, a more elegant version of the same could be found here