bamos / densenet.pytorch Goto Github PK

#!/usr/bin/env python3

import argparse
import os
import setproctitle
import shutil

import densenet
import torch
from torch import optim
from torch.autograd import Variable
from torch.nn import functional as F
from torch.utils.data import DataLoader
import torchvision
from torchvision import transforms


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batchSz', type=int, default=64)
    parser.add_argument('--nEpochs', type=int, default=300)
    parser.add_argument('--no-cuda', action='store_false')
    parser.add_argument('--save')
    parser.add_argument('--seed', type=int, default=1)
    parser.add_argument(
            '--opt', type=str, default='sgd',
            choices=('sgd', 'adam', 'rmsprop'))
    args = parser.parse_args()

    args.cuda = args.no_cuda and torch.cuda.is_available()
    if args.cuda:
        torch.cuda.manual_seed(args.seed)

    args.save = args.save or 'work/densenet.base'
    setproctitle.setproctitle(args.save)
    if os.path.exists(args.save):
        shutil.rmtree(args.save)
    os.makedirs(args.save, exist_ok=True)

    torch.manual_seed(args.seed)

    normMean = [0.49139968, 0.48215827, 0.44653124]
    normStd = [0.24703233, 0.24348505, 0.26158768]
    normTransform = transforms.Normalize(normMean, normStd)
    trainTransform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normTransform
    ])
    testTransform = transforms.Compose([
            transforms.ToTensor(),
            normTransform
    ])

    kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
    trainLoader = DataLoader(
            torchvision.datasets.CIFAR10(
                    root='cifar',
                    train=True,
                    download=True,
                    transform=trainTransform),
            batch_size=args.batchSz, shuffle=True, **kwargs)
    testLoader = DataLoader(
            torchvision.datasets.CIFAR10(
                    root='cifar',
                    train=False,
                    download=True,
                    transform=testTransform),
            batch_size=args.batchSz, shuffle=False, **kwargs)

    net = densenet.DenseNet(
            growthRate=12,
            depth=100,
            reduction=0.5,
            bottleneck=True,
            nClasses=10)

    print('  + Number of params: {}'.format(
            sum([p.data.nelement() for p in net.parameters()])))
    if args.cuda:
        net = net.cuda(1)

    if args.opt == 'sgd':
        optimizer = optim.SGD(
                net.parameters(), lr=1e-1, momentum=0.9, weight_decay=1e-4)
    elif args.opt == 'adam':
        optimizer = optim.Adam(net.parameters(), weight_decay=1e-4)
    elif args.opt == 'rmsprop':
        optimizer = optim.RMSprop(net.parameters(), weight_decay=1e-4)

    trainF = open(os.path.join(args.save, 'train.csv'), 'w')
    testF = open(os.path.join(args.save, 'test.csv'), 'w')

    for epoch in range(1, args.nEpochs + 1):
        adjust_opt(args.opt, optimizer, epoch)
        train(args, epoch, net, trainLoader, optimizer, trainF)
        test(args, epoch, net, testLoader, optimizer, testF)
        torch.save(net, os.path.join(args.save, 'latest.pth'))
        os.system('./plot.py {} &'.format(args.save))

    trainF.close()
    testF.close()


def train(args, epoch, net, trainLoader, optimizer, trainF):
    net.train()
    nProcessed = 0
    nTrain = len(trainLoader.dataset)
    for batch_idx, (data, target) in enumerate(trainLoader):
        if args.cuda:
            data, target = data.cuda(1), target.cuda(1)
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = net(data)
        loss = F.nll_loss(output, target)
        # make_graph.save('/tmp/t.dot', loss.creator); assert(False)
        loss.backward()
        optimizer.step()
        nProcessed += len(data)
        pred = output.data.max(1)[1]
        # get the index of the max log-probability
        incorrect = pred.ne(target.data).cpu().sum()
        err = 100.0 * incorrect / len(data)
        partialEpoch = epoch + batch_idx / len(trainLoader) - 1
        print(
                'Train Epoch: {:.2f} [{}/{}\t'
                '({:.0f}%)]\n'
                'Loss: {:.6f}\t' 'Error: {:.6f}'.format(
                        partialEpoch, nProcessed, nTrain,
                        100. * batch_idx / len(trainLoader),
                        loss.data[0], err))
        trainF.write('{},{},{}\n'.format(partialEpoch, loss.data[0], err))
        trainF.flush()


def test(args, epoch, net, testLoader, optimizer, testF):
    net.eval()
    test_loss = 0
    incorrect = 0
    for data, target in testLoader:
        if args.cuda:
            data, target = data.cuda(1), target.cuda(1)
        data, target = Variable(data, volatile=True), Variable(target)
        output = net(data)
        test_loss += F.nll_loss(output, target).data[0]
        pred = output.data.max(1)[1]
        # get the index of the max log-probability
        incorrect += pred.ne(target.data).cpu().sum()
    test_loss = test_loss
    test_loss /= len(testLoader)
    # loss function already averages over batch size
    nTotal = len(testLoader.dataset)
    err = 100.0 * incorrect / nTotal
    print()
    print(
            'Test set: Average loss: {:.4f}\n'
            'Error: {}/{} ({:.0f}%)\n'.format(
                    test_loss,
                    incorrect, nTotal, err))

    testF.write('{},{},{}\n'.format(epoch, test_loss, err))
    testF.flush()


def adjust_opt(optAlg, optimizer, epoch):
    if optAlg == 'sgd':
        if epoch < 150:
            lr = 1e-1
        elif epoch == 150:
            lr = 1e-2
        elif epoch == 225:
            lr = 1e-3
        else:
            return

        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

if __name__ == '__main__':
    main()