augment-ignore-dartslikelihood-resnet.py

""" Training augmented model """
import os
import torch
import torch.nn as nn
import numpy as np
from tensorboardX import SummaryWriter
from config import AugmentConfig
import utils
from models.augment_cnn import AugmentCNN
import copy
from models import MODEL_DICT


config = ResnetConfig()

device = torch.device("cuda")

# tensorboard
writer = SummaryWriter(log_dir=os.path.join(config.path, "tb"))
writer.add_text('config', config.as_markdown(), 0)

logger = utils.get_logger(os.path.join(config.path, "{}.log".format(config.name)))
config.print_params(logger.info)

class Architect():
    """ Compute gradients of alphas """
    def __init__(self, net, w_momentum, w_weight_decay):
        """
        Args:
            net
            w_momentum: weights momentum
        """
        self.net = net
        self.v_net = copy.deepcopy(net)
        self.w_momentum = w_momentum
        self.w_weight_decay = w_weight_decay

    def virtual_step(self, trn_X, trn_y, xi, w_optim, Likelihood, batch_size, step):
        """
        Compute unrolled weight w' (virtual step)

        Step process:
        1) forward
        2) calc loss
        3) compute gradient (by backprop)
        4) update gradient

        Args:
            xi: learning rate for virtual gradient step (same as weights lr)
            w_optim: weights optimizer
        """
        # forward & calc loss
        dataIndex = len(trn_y)+step*batch_size
        ignore_crit = nn.CrossEntropyLoss(reduction='none').cuda()       
        # forward
        logits = self.net(trn_X)
         
        # sigmoid loss
        first = torch.sigmoid(Likelihood[step*batch_size:dataIndex])
        second = ignore_crit(logits, trn_y)
        lossup = torch.dot(first,second )
        lossdiv =(torch.sigmoid(Likelihood[step*batch_size:dataIndex]).sum())
        loss = lossup/lossdiv
        
#         loss = torch.dot(torch.sigmoid(Likelihood[step*batch_size:dataIndex]), ignore_crit(logits, trn_y))/(torch.sigmoid(Likelihood[step*batch_size:dataIndex]).sum())
        
        # compute gradient of train loss towards likelihhod
        loss.backward()

        # do virtual step (update gradient)
        # below operations do not need gradient tracking
        with torch.no_grad():
            # dict key is not the value, but the pointer. So original network weight have to
            # be iterated also.
            for w, vw in zip(self.net.parameters(), self.v_net.parameters()):
                m = w_optim.state[w].get('momentum_buffer', 0.) * self.w_momentum
                
                if w.grad is not None:
                    vw.copy_(w - xi * (m + w.grad + self.w_weight_decay*w))


    def unrolled_backward(self, trn_X, trn_y, val_X, val_y, xi, w_optim, Likelihood, Likelihood_optim, batch_size, step):
        """ Compute unrolled loss and backward its gradients
        Args:
            xi: learning rate for virtual gradient step (same as net lr)
            w_optim: weights optimizer - for virtual step
        """
        crit = nn.CrossEntropyLoss().to(device)
        
        # do virtual step (calc w`)
        self.virtual_step(trn_X, trn_y, xi, w_optim, Likelihood, batch_size, step)
        
        # calc val prediction
        logits = self.v_net(val_X)   
        # calc unrolled validation loss
        loss = crit(logits, val_y) # L_val(w`)
        
        # compute gradient of validation loss towards weights
        v_weights = tuple(self.v_net.parameters())
        # some weights not used return none
        dw = torch.autograd.grad(loss, v_weights, allow_unused=True)
           
        hessian = self.compute_hessian(dw, trn_X, trn_y, Likelihood, batch_size, step)

        
        # validation precision   
        vprec1, vprec5 = utils.accuracy(logits, val_y, topk=(1, 5))
        
        Likelihood_optim.zero_grad()
        # update final gradient = - xi*hessian
#         with torch.no_grad():
#             for likelihood, h in zip(Likelihood, hessian):
#                 print(len(hessian))
#                 likelihood.grad = - xi*h
        with torch.no_grad():
            Likelihood.grad = - xi*hessian[0]         
        Likelihood_optim.step()
        return Likelihood, Likelihood_optim, loss, vprec1, vprec5

    def compute_hessian(self, dw, trn_X, trn_y, Likelihood, batch_size, step):
        """
        dw = dw` { L_val(w`, alpha) }
        w+ = w + eps * dw
        w- = w - eps * dw
        hessian = (dalpha { L_trn(w+, alpha) } - dalpha { L_trn(w-, alpha) }) / (2*eps)
        eps = 0.01 / ||dw||
        """
        norm = torch.cat([w.view(-1) for w in dw if w != None]).norm()
        
        eps = 0.01 / norm
        
        # w+ = w + eps*dw`
        with torch.no_grad():
            for p, d in zip(self.net.parameters(), dw):
                if d != None:
                    p += eps * d
        
        
        # forward & calc loss
        dataIndex = len(trn_y)+step*batch_size
        ignore_crit = nn.CrossEntropyLoss(reduction='none').cuda()       
        # forward
        logits = self.net(trn_X)       
        # sigmoid loss
        loss = torch.dot(torch.sigmoid(Likelihood[step*batch_size:dataIndex]), ignore_crit(logits, trn_y))/(torch.sigmoid(Likelihood[step*batch_size:dataIndex]).sum())
        
        
        dalpha_pos = torch.autograd.grad(loss, Likelihood) # dalpha { L_trn(w+) }

        # w- = w - eps*dw`
        with torch.no_grad():
            for p, d in zip(self.net.parameters(), dw):
                if d != None:
                    p -= 2. * eps * d
        # forward
        logits = self.net(trn_X)       
        # sigmoid loss
        loss = torch.dot(torch.sigmoid(Likelihood[step*batch_size:dataIndex]), ignore_crit(logits, trn_y))/(torch.sigmoid(Likelihood[step*batch_size:dataIndex]).sum())
        dalpha_neg = torch.autograd.grad(loss, Likelihood) # dalpha { L_trn(w-) }

        # recover w
        with torch.no_grad():
            for p, d in zip(self.net.parameters(), dw):
                if d != None:
                    p += eps * d

        hessian = [(p-n) / 2.*eps for p, n in zip(dalpha_pos, dalpha_neg)]
#         hessian = [(p-n) / (2.*eps) for p, n in zip(dalpha_pos, dalpha_neg)]
        return hessian

    
def main():
    logger.info("Logger is set - training start")

    # set default gpu device id
    torch.cuda.set_device(config.gpus[0])

    # set seed
    np.random.seed(config.seed)
    torch.manual_seed(config.seed)
    torch.cuda.manual_seed_all(config.seed)

    torch.backends.cudnn.benchmark = True

    # get data with meta info
    input_size, input_channels, n_classes, train_val_data, test_data = utils.get_data(
        config.dataset, config.data_path, config.cutout_length, validation=True)

    criterion = nn.CrossEntropyLoss().to(device)
    use_aux = config.aux_weight > 0.
    model = MODEL_DICT['resnet18']().to(device)
#     model = nn.DataParallel(model, device_ids=config.gpus).to(device)

    # model size
    mb_params = utils.param_size(model)
    logger.info("Model size = {:.3f} MB".format(mb_params))

    # weights optimizer with SGD
    optimizer = torch.optim.SGD(model.parameters(), config.lr, momentum=config.momentum,
                                weight_decay=config.weight_decay)
    
    
    n_train = len(train_val_data)
    split = n_train // 2
    indices = list(range(n_train))
    
    # each train data is endowed with a weight
    Likelihood = torch.nn.Parameter(torch.ones(len(indices[:split])).cuda(),requires_grad=True).cuda()
#     Likelihood_optim = torch.optim.SGD({Likelihood}, config.lr)
    Likelihood_optim = torch.optim.Adam({Likelihood}, config.alpha_lr, betas=(0.5, 0.999))
    
    
    # data split
    train_data = torch.utils.data.Subset(train_val_data, indices[:split])
    valid_data = torch.utils.data.Subset(train_val_data, indices[split:])
    
    
    train_loader = torch.utils.data.DataLoader(train_data,
                                               batch_size=config.batch_size,
                                               shuffle=False,
                                               num_workers=config.workers,
                                               pin_memory=False)
    valid_loader = torch.utils.data.DataLoader(valid_data,
                                               batch_size=config.batch_size,
                                               shuffle=False,
                                               num_workers=config.workers,
                                               pin_memory=False)
        
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, config.epochs)
    architect = Architect(model, 0.9, 3e-4)
    
    best_top1 = 0.
    # training loop
    for epoch in range(config.epochs):
        lr_scheduler.step()
        lr = lr_scheduler.get_lr()[0]
        drop_prob = config.drop_path_prob * epoch / config.epochs
#         model.drop_path_prob(drop_prob)

        # training
        train(train_loader, valid_loader, model, architect, optimizer, criterion, lr, epoch, Likelihood, Likelihood_optim, config.batch_size)

        # validation
        cur_step = (epoch+1) * len(train_loader)
        top1 = validate(valid_loader, model, criterion, epoch, cur_step)

        # save
        if best_top1 < top1:
            best_top1 = top1
            is_best = True
        else:
            is_best = False
        utils.save_checkpoint(model, config.path, is_best)

        print("")

    logger.info("Final best Prec@1 = {:.4%}".format(best_top1))


def train(train_loader, valid_loader, model, architect, optimizer, criterion, lr, epoch, Likelihood, Likelihood_optim, batch_size):
    top1 = utils.AverageMeter()
    top5 = utils.AverageMeter()
    losses = utils.AverageMeter()
    standard_losses = utils.AverageMeter()
    valid_losses = utils.AverageMeter()

    cur_step = epoch*len(train_loader)
    cur_lr = optimizer.param_groups[0]['lr']
    logger.info("Epoch {} LR {}".format(epoch, cur_lr))
    writer.add_scalar('train/lr', cur_lr, cur_step)

    model.train()
    
    for step, ((trn_X, trn_y), (val_X, val_y)) in enumerate(zip(train_loader, valid_loader)):
        trn_X, trn_y = trn_X.to(device, non_blocking=True), trn_y.to(device, non_blocking=True)
        val_X, val_y = val_X.to(device, non_blocking=True), val_y.to(device, non_blocking=True)
        N = trn_X.size(0)
        M = val_X.size(0)

        # phase 2. Likelihood step (Likelihood)
        Likelihood_optim.zero_grad()
        Likelihood, Likelihood_optim, valid_loss, vprec1, vprec5= architect.unrolled_backward(trn_X, trn_y, val_X, val_y, lr, optimizer, Likelihood, Likelihood_optim, batch_size, step)
            
        print(Likelihood)
        print(Likelihood.sum())
            
        # phase 1. network weight step (w)    
        optimizer.zero_grad()
        logits = model(trn_X)      
        
        ignore_crit = nn.CrossEntropyLoss(reduction='none').to(device)
        dataIndex = len(trn_y)+step*batch_size
        loss = torch.dot(torch.sigmoid(Likelihood[step*batch_size:dataIndex]), ignore_crit(logits, trn_y))
        loss = loss/(torch.sigmoid(Likelihood[step*batch_size:dataIndex]).sum())
        '''
        if config.aux_weight > 0.:
            loss += config.aux_weight * criterion(aux_logits, y)
        '''
        loss.backward()
        # gradient clipping
        nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
        # update network weight on train data
        optimizer.step()     
        
        #compare normal loss without weighted
        standard_loss = criterion(logits, trn_y)
            
        prec1, prec5 = utils.accuracy(logits, trn_y, topk=(1, 5))
        losses.update(loss.item(), N)
        standard_losses.update(standard_loss.item(), N)
        valid_losses.update(valid_loss.item(), M)
        top1.update(prec1.item(), N)
        top5.update(prec5.item(), N)
    
        if step % config.print_freq == 0 or step == len(train_loader)-1:
            logger.info(
                "Train: [{:3d}/{}] Step {:03d}/{:03d} Loss {losses.avg:.3f} standard Loss {slosses.avg:.3f} Valid Loss {vlosses.avg:.3f}"
                " Prec@(1,5) ({top1.avg:.1%}, {top5.avg:.1%})".format(
                    epoch+1, config.epochs, step, len(train_loader)-1, losses=losses, slosses=standard_losses, vlosses=valid_losses,
                    top1=top1, top5=top5))
        

        writer.add_scalar('train/loss', loss.item(), cur_step)
        writer.add_scalar('train/top1', prec1.item(), cur_step)
        writer.add_scalar('train/top5', prec5.item(), cur_step)
        writer.add_scalar('val/loss', valid_loss.item(), cur_step)
        writer.add_scalar('val/top1', vprec1.item(), cur_step)
        writer.add_scalar('val/top5', vprec5.item(), cur_step)
        cur_step += 1

    logger.info("Train: [{:3d}/{}] Final Prec@1 {:.4%}".format(epoch+1, config.epochs, top1.avg))


def validate(valid_loader, model, criterion, epoch, cur_step):
    top1 = utils.AverageMeter()
    top5 = utils.AverageMeter()
    losses = utils.AverageMeter()

    model.eval()

    with torch.no_grad():
        for step,(X, y) in enumerate(valid_loader):
            X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
            N = X.size(0) 

            logits = model(X)
            loss = criterion(logits, y)
            
            prec1, prec5 = utils.accuracy(logits, y, topk=(1, 5))
            losses.update(loss.item(), N)
            top1.update(prec1.item(), N)
            top5.update(prec5.item(), N)

            if step % config.print_freq == 0 or step == len(valid_loader)-1:
                logger.info(
                    "Test: [{:3d}/{}] Step {:03d}/{:03d} Loss {losses.avg:.3f} "
                    "Prec@(1,5) ({top1.avg:.1%}, {top5.avg:.1%})".format(
                        epoch+1, config.epochs, step, len(valid_loader)-1, losses=losses,
                        top1=top1, top5=top5))

    writer.add_scalar('test/loss', losses.avg, cur_step)
    writer.add_scalar('test/top1', top1.avg, cur_step)
    writer.add_scalar('test/top5', top5.avg, cur_step)

    logger.info("Test: [{:3d}/{}] Final Prec@1 {:.4%}".format(epoch+1, config.epochs, top1.avg))

    return top1.avg


if __name__ == "__main__":
    main()