train.py

# train.py
#!/usr/bin/env	python3

""" train network using pytorch

author baiyu
"""

import os
import sys
import argparse
from datetime import datetime

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

from torch.utils.data import DataLoader
#from dataset import *
from torch.autograd import Variable

from tensorboardX import SummaryWriter

from conf import settings
from utils import get_network, get_training_dataloader, get_test_dataloader, WarmUpLR
from attack_config import config
from training.train import eval_one_epoch,eval_one_epoch2, train_one_epoch

def train(epoch):

    net.train()
    for batch_index, (images, labels) in enumerate(cifar10_training_loader):
        if epoch <= args.warm:
            warmup_scheduler.step()

        images = Variable(images)
        labels = Variable(labels)
        labels = labels.cuda()
        images = images.cuda()

        optimizer.zero_grad()
        outputs = net(images)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        if args.net == 'resnet18':
            with torch.no_grad():
                net.fc.weight.div_(torch.norm(net.fc.weight))

        n_iter = (epoch - 1) * len(cifar10_training_loader) + batch_index + 1

        last_layer = list(net.children())[-1]
        # for name, para in last_layer.named_parameters():
        #     if 'weight' in name:
        #         writer.add_scalar('LastLayerGradients/grad_norm2_weights', para.grad.norm(), n_iter)
        #     if 'bias' in name:
        #         writer.add_scalar('LastLayerGradients/grad_norm2_bias', para.grad.norm(), n_iter)

        print('Training Epoch: {epoch} [{trained_samples}/{total_samples}]\tLoss: {:0.4f}\tLR: {:0.6f}'.format(
            loss.item(),
            optimizer.param_groups[0]['lr'],
            epoch=epoch,
            trained_samples=batch_index * args.b + len(images),
            total_samples=len(cifar10_training_loader.dataset)
        ))

        #update training loss for each iteration
        writer.add_scalar('Train/loss', loss.item(), n_iter)

    for name, param in net.named_parameters():
        layer, attr = os.path.splitext(name)
        attr = attr[1:]
        writer.add_histogram("{}/{}".format(layer, attr), param, epoch)

def eval_training(epoch):
    net.eval()

    test_loss = 0.0 # cost function error
    correct = 0.0

    for (images, labels) in cifar10_test_loader:
        images = Variable(images)
        labels = Variable(labels)

        images = images.cuda()
        labels = labels.cuda()

        outputs = net(images)
        loss = loss_function(outputs, labels)
        test_loss += loss.item()
        _, preds = outputs.max(1)
        correct += preds.eq(labels).sum()

    print('Test set: Average loss: {:.4f}, Accuracy: {:.4f}'.format(
        test_loss / len(cifar10_test_loader.dataset),
        correct.float() / len(cifar10_test_loader.dataset)
    ))
    print()

    #add informations to tensorboard
    writer.add_scalar('Test/Average loss', test_loss / len(cifar10_test_loader.dataset), epoch)
    writer.add_scalar('Test/Accuracy', correct.float() / len(cifar10_test_loader.dataset), epoch)

    return correct.float() / len(cifar10_test_loader.dataset)

def eval_training2(epoch):
    net.eval()

    test_loss = 0.0 # cost function error
    correct = 0.0

    for (images, labels) in cifar10_test_loader:
        images = Variable(images)
        labels = Variable(labels)

        images = images.cuda()
        labels = labels.cuda()

        outputs = net(images)
        loss = loss_function(outputs, labels)
        test_loss += loss.item()
        _, preds = outputs.max(1)
        correct += preds.eq(labels).sum()

    print('Test set: Average loss: {:.4f}, Accuracy: {:.4f}'.format(
        test_loss / len(cifar10_test_loader.dataset),
        correct.float() / len(cifar10_test_loader.dataset)
    ))
    print()

    #add informations to tensorboard
    writer.add_scalar('Test/Average loss', test_loss / len(cifar10_test_loader.dataset), epoch)
    writer.add_scalar('Test/Accuracy', correct.float() / len(cifar10_test_loader.dataset), epoch)

    return correct.float() / len(cifar10_test_loader.dataset)

if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('-net', type=str, default='sir_resnet34', help='net type')
    parser.add_argument('-gpu', type=bool, default=True, help='use gpu or not')
    parser.add_argument('-w', type=int, default=2, help='number of workers for dataloader')
    parser.add_argument('-b', type=int, default=128, help='batch size for dataloader')
    parser.add_argument('-s', type=bool, default=True, help='whether shuffle the dataset')
    parser.add_argument('-warm', type=int, default=1, help='warm up training phase')
    parser.add_argument('-lr', type=float, default=0.1, help='initial learning rate')
    parser.add_argument('-attack',type=bool,default=True)
    parser.add_argument('--wd', type=float, default=0.1)
    parser.add_argument('--it', type=int, default=1)
    parser.add_argument('-train',type=str,default='pgd3')
    args = parser.parse_args()

    net = get_network(args, use_gpu=args.gpu)
        
    #data preprocessing:
    cifar10_training_loader = get_training_dataloader(
        settings.CIFAR100_TRAIN_MEAN,
        settings.CIFAR100_TRAIN_STD,
        num_workers=args.w,
        batch_size=args.b,
        shuffle=args.s
    )
    
    cifar10_test_loader = get_test_dataloader(
        settings.CIFAR100_TRAIN_MEAN,
        settings.CIFAR100_TRAIN_STD,
        num_workers=args.w,
        batch_size=args.b,
        shuffle=args.s
    )
    
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
    train_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=settings.MILESTONES2, gamma=0.1) #learning rate decay
    iter_per_epoch = len(cifar10_training_loader)
    if args.warm>0:
        warmup_scheduler = WarmUpLR(optimizer, iter_per_epoch * args.warm)
    checkpoint_path = os.path.join(settings.CHECKPOINT_PATH, 'CIFAR10-2', args.net, args.train+' '+str(args.wd)+'-'+str(args.it)+'-'+settings.TIME_NOW)

    #use tensorboard
    if not os.path.exists(settings.LOG_DIR):
        os.mkdir(settings.LOG_DIR)
    writer = SummaryWriter(log_dir=os.path.join(
            settings.LOG_DIR, args.net, settings.TIME_NOW))
    input_tensor = torch.Tensor(12, 3, 32, 32).cuda()
    # writer.add_graph(net, Variable(input_tensor, requires_grad=True))

    #create checkpoint folder to save model
    if not os.path.exists(checkpoint_path):
        os.makedirs(checkpoint_path)
    checkpoint_path = os.path.join(checkpoint_path, '{net}-{epoch}-{type}.pth')

    best_acc = 0.0
    
    for epoch in range(1, settings.EPOCH2):
        if epoch > args.warm:
            train_scheduler.step(epoch)
        # print('Epoch {}'.format(epoch))

        # train(epoch)
        if args.warm>0 and epoch<=args.warm:
            train(epoch)
        elif args.train == 'pgd5':
            trainattack = config.create_evaluation_attack_method4(torch.device('cuda:{}'.format(0)))
            train_one_epoch(net = net, batch_generator = cifar10_training_loader, optimizer= optimizer, criterion= loss_function, DEVICE = torch.device('cuda:{}'.format(0)), descrip_str = str(epoch)+ 'Training',
             AttackMethod = trainattack)
        elif args.train == 'pgd3':
            trainattack = config.create_attack_method(torch.device('cuda:{}'.format(0)))
            train_one_epoch(net = net, batch_generator = cifar10_training_loader, optimizer= optimizer, criterion= loss_function, DEVICE = torch.device('cuda:{}'.format(0)), descrip_str = str(epoch)+ 'Training',
             AttackMethod = trainattack)
        elif args.train == 'pgd10':
            trainattack = config.create_attack_method10(torch.device('cuda:{}'.format(0)))
            train_one_epoch(net = net, batch_generator = cifar10_training_loader, optimizer= optimizer, criterion= loss_function, DEVICE = torch.device('cuda:{}'.format(0)), descrip_str = str(epoch)+ 'Training',
             AttackMethod = trainattack)
        elif args.train == 'pgd1':
            trainattack = config.create_evaluation_attack_method(torch.device('cuda:{}'.format(0)))
            train_one_epoch(net = net, batch_generator = cifar10_training_loader, optimizer= optimizer, criterion= loss_function, DEVICE = torch.device('cuda:{}'.format(0)), descrip_str = str(epoch)+ 'Training',
             AttackMethod = trainattack)
        else:
            # train_one_epoch(net = net, batch_generator = cifar10_training_loader, optimizer= optimizer, criterion= loss_function, DEVICE = torch.device('cuda:{}'.format(0)), descrip_str = str(epoch)+ 'Training',
            #  AttackMethod = None)
            trainattack = None
            train_one_epoch(net = net, batch_generator = cifar10_training_loader, optimizer= optimizer, criterion= loss_function, DEVICE = torch.device('cuda:{}'.format(0)), descrip_str = str(epoch)+ 'Training',
             AttackMethod = trainattack)
        print('Learning_rate:', optimizer.param_groups[0]['lr'])
        if args.attack:
            AttackMethod = config.create_evaluation_attack_method(torch.device('cuda:{}'.format(0)))
        acc = eval_training(epoch)
        clean_acc, adv_acc = eval_one_epoch(net, cifar10_test_loader, torch.device('cuda:{}'.format(0)), AttackMethod)
        print('clean acc -- {}     adv acc -- {} in training mode'.format(clean_acc, adv_acc))
        # acc = eval_training(epoch)
        
        # clean_acc, adv_acc_t = eval_one_epoch(net, cifar10_test_loader, torch.device('cuda:{}'.format(0)), AttackMethod)
        # print('clean acc -- {}     adv acc -- {}'.format(clean_acc, adv_acc_t))
        
        #start to save best performance model after learning rate decay to 0.01 
        if epoch > settings.MILESTONES2[1] and best_acc < adv_acc:
            torch.save(net.state_dict(), checkpoint_path.format(net=args.net, epoch=epoch, type='best'))
            best_acc = adv_acc
            continue

        if not epoch % 30 or epoch>170:
            torch.save(net.state_dict(), checkpoint_path.format(net=args.net, epoch=epoch, type='regular'))
    print(args.net,"The best adversarial acc is {}".format(best_acc))
    writer.close()