train.py

##### System library #####
import os
import os.path as osp
from os.path import exists
import argparse
import json
import logging
import time
import numpy as np
import shutil
##### pytorch library #####
import torch
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
import torch.nn.functional as F
##### My own library #####
import data.seg_transforms as dt
from data.Seg_dataset import SegList
from utils.logger import Logger
from models.net_builder import net_builder
from utils.loss import loss_builder
from utils.utils import compute_average_dice, AverageMeter, save_checkpoint, count_param, target_seg2target_cls
from utils.utils import aic_fundus_lesion_segmentation, aic_fundus_lesion_classification, compute_segment_score, \
    compute_single_segment_score

# logger vis
FORMAT = "[%(asctime)-15s %(filename)s:%(lineno)d %(funcName)s] %(message)s"
logging.basicConfig(format=FORMAT)
logger_vis = logging.getLogger(__name__)
logger_vis.setLevel(logging.DEBUG)

###### train #######
def adjust_learning_rate(args, optimizer, epoch):
    """
    Sets the learning rate to the initial LR decayed by 10 every 30 epochs(step = 30)
    """
    if args.lr_mode == 'step':
        lr = args.lr * (0.1 ** (epoch // args.step))
    elif args.lr_mode == 'poly':
        lr = args.lr * (1 - epoch / args.epochs) ** 0.9
    else:
        raise ValueError('Unknown lr mode {}'.format(args.lr_mode))

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


def train(args, train_loader, model, criterion, optimizer, epoch, print_freq=10):
    # set the AverageMeter
    batch_time = AverageMeter()
    losses = AverageMeter()
    dice = AverageMeter()
    Dice_1 = AverageMeter()
    Dice_2 = AverageMeter()
    Dice_3 = AverageMeter()
    # switch to train mode
    model.train()
    end = time.time()
    correct, total = 0, 0

    for i, (input, target,id, patient) in enumerate(train_loader):

        target_seg = target.numpy()
        target_cls = target_seg2target_cls(target_seg).cuda()  # transfer seg target to cls target
        # Variable
        input_var = Variable(input).cuda()
        target_var_seg = Variable(target).cuda()
        target_var_cls = Variable(target_cls).cuda()
        id_1 = id[::2]
        id_2 = id[1::2]
        patient_1 = patient[::2]
        patient_2 = patient[1::2]
        target_var_seg_1 = target_var_seg[::2]
        target_var_seg_2 = target_var_seg[1::2]
        target_var_cls_1 = target_var_cls[::2]
        target_var_cls_2 = target_var_cls[1::2]
        # forward

        input_var_1 = input_var[::2]
        input_var_2 = input_var[1::2]
        output_seg_1, output_cls_1, cls_logits_1, seg_logits_1 = model(input_var_1)
        output_seg_2, output_cls_2, cls_logits_2, seg_logits_2 = model(input_var_2)

        output_seg = torch.cat([output_seg_1, output_seg_2], dim=0)
        output_cls = torch.cat([output_cls_1, output_cls_2], dim=0)

        loss_seg_1 = criterion[0](output_seg_1, target_var_seg_1)
        loss_seg_2 = criterion[0](output_seg_2, target_var_seg_2)
        loss_seg = (loss_seg_1 + loss_seg_2) / 2

        loss_cls_1 = criterion[2](output_cls_1, target_var_cls_1)
        loss_cls_2 = criterion[2](output_cls_2, target_var_cls_2)
        loss_cls = (loss_cls_1 + loss_cls_2) / 2
        c = 0
        cls_con_loss = 0
        seg_con_loss = 0
        loss_super = loss_seg + loss_cls

        loss_back = torch.nn.KLDivLoss().cuda()
        back_region = input_var[:, :, 496 - 70:496 - 30, :]
        background = torch.zeros_like(input_var)
        yushu = input_var.shape[2] % 40
        background[:, :, :yushu, :] = back_region[:, :, :yushu, :]
        for k in range(int(496 / 40)):
            background[:, :, yushu + k * 40:yushu + (k + 1) * 40, :] = back_region

        seg_region = torch.mul(F.softmax(1000*output_seg, dim=1),input_var)
        loss_background = loss_back(seg_region, background.cuda())

        loss_cls_con_L2 = torch.nn.MSELoss(reduction='mean')
        loss_seg_con_L1 = torch.nn.L1Loss(reduction='mean')
        loss_inter = 0
        for j in range(input_var_1.shape[0]):
            for k in range(input_var_1.shape[0]):
                if patient_1[j] == patient_2[k] and abs(id_1[j] - id_2[k]) < 3:
                    c += 1
                    cls_con_loss = cls_con_loss + loss_cls_con_L2(cls_logits_1[j], cls_logits_2[k])

                    seg_con_loss = seg_con_loss + loss_seg_con_L1(seg_logits_1[j], seg_logits_2[k])

                    print('find %s, id:%d, %d' % (patient_1[j], id_1[j], id_2[k]))
        if c > 0:
            loss_inter = seg_con_loss / c + cls_con_loss / c

        loss = loss_super +  loss_background + loss_inter

        losses.update(loss.data, input.size(0))
        # metric dice for seg
        _, pred_seg = torch.max(output_seg, 1)
        pred_seg = pred_seg.cpu().data.numpy()
        label_seg = target_var_seg.cpu().data.numpy()
        dice_score, dice_1, dice_2, dice_3 = compute_average_dice(pred_seg.flatten(), label_seg.flatten())
        # metric acc for cls
        pred_cls = (output_cls > 0.5)
        total += target_var_cls.size(0) * 3
#        correct += pred_cls.eq(target_var_cls.byte()).sum().item()
        # update dice
        dice.update(dice_score)
        Dice_1.update(dice_1)
        Dice_2.update(dice_2)
        Dice_3.update(dice_3)
        # backwards
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        # logger vis
        if i % print_freq == 0:
            logger_vis.info('Epoch: [{0}][{1}/{2}]\t'
                            'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                            'Dice {dice.val:.4f} ({dice.avg:.4f})\t'
                            'Dice_1 {dice_1.val:.6f} ({dice_1.avg:.4f})\t'
                            'Dice_2 {dice_2.val:.6f} ({dice_2.avg:.4f})\t'
                            'Dice_3 {dice_3.val:.6f} ({dice_3.avg:.4f})\t'.format(
                epoch, i, len(train_loader), batch_time=batch_time, dice=dice, dice_1=Dice_1, dice_2=Dice_2, dice_3=Dice_3, loss=losses))
            print('loss: %.3f (%.4f)' %(losses.val,losses.avg))
            #'Loss {loss.val:.6f} ({liss.avg:.4f})\t'
        #break
    return losses.avg, dice.avg, Dice_1.avg, Dice_2.avg, Dice_3.avg


def train_seg(args, result_path, logger):
    for k, v in args.__dict__.items():
        print(k, ':', v)
    # load the net
    net = net_builder(args.name, args.model_path, args.pretrained)
    model = torch.nn.DataParallel(net).cuda()
    param = count_param(model)
    print('###################################')
    print('Model #%s# parameters: %.2f M' % (args.name, param / 1e6))
    # set the loss criterion
    criterion = loss_builder(args.loss)
    # Data loading code
    info = json.load(open(osp.join(args.list_dir, 'info.json'), 'r'))
    normalize = dt.Normalize(mean=info['mean'], std=info['std'])
    # data transforms
    t = []
    if args.resize:
        t.append(dt.Resize(args.resize))
    if args.random_rotate > 0:
        t.append(dt.RandomRotate(args.random_rotate))
    if args.random_scale > 0:
        t.append(dt.RandomScale(args.random_scale))
    if args.crop_size:
        t.append(dt.RandomCrop(args.crop_size))
    t.extend([dt.Label_Transform(),
              dt.RandomHorizontalFlip(),
              dt.ToTensor(),
              normalize])
    # dataset = SegList(args.data_dir, 'train', dt.Compose(t),list_dir=args.list_dir)
    train_loader = torch.utils.data.DataLoader(
        SegList(args.data_dir, 'train', dt.Compose(t), list_dir=args.list_dir), batch_size=args.batch_size,
        shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True)

    # define loss function (criterion) and pptimizer
    if args.optimizer == 'SGD':  # SGD optimizer
        optimizer = torch.optim.SGD(net.parameters(),
                                    args.lr,
                                    momentum=args.momentum,
                                    weight_decay=args.weight_decay)
    elif args.optimizer == 'Adam':  # Adam optimizer
        optimizer = torch.optim.Adam(net.parameters(),
                                     args.lr,
                                     betas=(0.9, 0.99),
                                     weight_decay=args.weight_decay)

    cudnn.benchmark = True
    best_dice = 0
    start_epoch = 0

    # load the pretrained model
    if args.model_path:
        print("=> loading pretrained model '{}'".format(args.model_path))
        checkpoint = torch.load(args.model_path)
        model.load_state_dict(checkpoint['state_dict'])
    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            start_epoch = checkpoint['epoch']
            best_dice = checkpoint['best_dice']
            dice_epoch = checkpoint['dice_epoch']
            model.load_state_dict(checkpoint['state_dict'])
            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    # main training
    for epoch in range(start_epoch, args.epochs):
        lr = adjust_learning_rate(args, optimizer, epoch)
        logger_vis.info('Epoch: [{0}]\tlr {1:.06f}'.format(epoch, lr))
        # train for one epoch
        # loss, dice_train, dice_1, dice_2 = train(args, train_loader, model, criterion, optimizer, epoch)
        loss, dice_train, dice_1, dice_2, dice_3 = train(args, train_loader, model, criterion, optimizer, epoch)
        # evaluate on validation set
        dice_val, dice_11, dice_22, dice_33, dice_list, auc, auc_1, auc_2, auc_3 = val_seg(args, model)
        # dice_val, dice_11, dice_22, dice_list, auc, auc_1, auc_2 = val_seg(args, model)
        # save best checkpoints
        is_best = dice_val > best_dice
        best_dice = max(dice_val, best_dice)
        checkpoint_dir = osp.join(result_path, 'checkpoint')
        if not exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        checkpoint_latest = checkpoint_dir + '/checkpoint_latest.pth.tar'
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'dice_epoch': dice_val,
            'best_dice': best_dice,
        }, is_best, checkpoint_dir, filename=checkpoint_latest)
        if args.save_every_checkpoint:
            if (epoch + 1) % 1 == 0:
                history_path = checkpoint_dir + '/checkpoint_{:03d}.pth.tar'.format(epoch + 1)
                shutil.copyfile(checkpoint_latest, history_path)

        # logger.append([epoch, dice_train, dice_val, auc, dice_1, dice_11, dice_2, dice_22, auc_1, auc_2])
        logger.append([epoch, dice_train, dice_val, auc, dice_11, dice_22, dice_33, auc_1, auc_2, auc_3])


####### validation ###########

def val(args, eval_data_loader, model):
    model.eval()
    batch_time = AverageMeter()
    dice = AverageMeter()
    end = time.time()
    dice_list = []
    Dice_1 = AverageMeter()
    Dice_2 = AverageMeter()
    Dice_3 = AverageMeter()
    ret_segmentation = []

    for iter, (image, label, _) in enumerate(eval_data_loader):
        # batchsize = 1 ,so squeeze dim 1
        image = image.squeeze(dim=0)
        label = label.squeeze(dim=0)

        target_seg = label.numpy()
        target_cls = target_seg2target_cls(target_seg)

        with torch.no_grad():
            # batch test for memory reduce
            batch = 16
            pred_seg = torch.zeros(image.shape[0], image.shape[2], image.shape[3])
            pred_cls = torch.zeros(image.shape[0], 3)
            for i in range(0, image.shape[0], batch):
                start_id = i
                end_id = i + batch
                if end_id > image.shape[0]:
                    end_id = image.shape[0]
                image_batch = image[start_id:end_id, :, :, :]
                image_var = Variable(image_batch).cuda()
                # model forward
                output_seg, output_cls, _, _  = model(image_var)
                _, pred_batch = torch.max(output_seg, 1)
                pred_seg[start_id:end_id, :, :] = pred_batch.cpu().data
                pred_cls[start_id:end_id, :] = output_cls.cpu().data
            # merice dice for seg
            pred_seg = pred_seg.numpy().astype('uint8')
            batch_time.update(time.time() - end)
            label_seg = label.numpy().astype('uint8')
            ret = aic_fundus_lesion_segmentation(label_seg, pred_seg)
            ret_segmentation.append(ret)
            dice_score = compute_single_segment_score(ret)
            dice_list.append(dice_score)
            dice.update(dice_score)
            Dice_1.update(ret[1])
            Dice_2.update(ret[2])
            Dice_3.update(ret[3])
            # metric auc for cls
            ground_truth = target_cls.numpy().astype('float32')
            prediction = pred_cls.numpy().astype('float32')
            if iter == 0:
                detection_ref_all = ground_truth
                detection_pre_all = prediction
            else:
                detection_ref_all = np.concatenate((detection_ref_all, ground_truth), axis=0)
                detection_pre_all = np.concatenate((detection_pre_all, prediction), axis=0)

        end = time.time()
        logger_vis.info('Eval: [{0}/{1}]\t'
                        'Dice {dice.val:.3f} ({dice.avg:.3f})\t'
                        'Dice_1 {dice_1.val:.3f} ({dice_1.avg:.3f})\t'
                        'Dice_2 {dice_2.val:.3f} ({dice_2.avg:.3f})\t'
                        'Dice_3 {dice_3.val:.3f} ({dice_3.avg:.3f})'
                        'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                        .format(iter, len(eval_data_loader), dice=dice, dice_1=Dice_1, dice_2=Dice_2, dice_3=Dice_3,
                                batch_time=batch_time))
        # break
    # compute average dice for seg
    final_seg, seg_1, seg_2, seg_3 = compute_segment_score(ret_segmentation)
    print('### Seg ###')
    print('Final Seg Score:{}'.format(final_seg))
    print('Final Seg_1 Score:{}'.format(seg_1))
    print('Final Seg_2 Score:{}'.format(seg_2))
    print('Final Seg_3 Score:{}'.format(seg_3))
    # compute average auc for cls
    ret_detection = aic_fundus_lesion_classification(detection_ref_all, detection_pre_all,
                                                     num_samples=len(eval_data_loader) * 128)
    auc = np.array(ret_detection).mean()
    print('AUC :', auc)
    auc_1 = ret_detection[0]
    auc_2 = ret_detection[1]
    auc_3 = ret_detection[2]

    return final_seg, seg_1, seg_2, seg_3, dice_list, auc, auc_1, auc_2, auc_3


def val_seg(args, model):
    info = json.load(open(osp.join(args.list_dir, 'info.json'), 'r'))
    normalize = dt.Normalize(mean=info['mean'], std=info['std'])

    t = []
    if args.resize:
        t.append(dt.Resize(args.resize))
    if args.crop_size:
        t.append(dt.RandomCrop(args.crop_size))
    t.extend([dt.Label_Transform(),
              dt.ToTensor(),
              normalize])

    dataset = SegList(args.data_dir, 'val', dt.Compose(t), list_dir=args.list_dir)
    val_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False,
                                             num_workers=args.workers, pin_memory=False)

    cudnn.benchmark = True
    # dice_avg, dice_1, dice_2, dice_list, auc, auc_1, auc_2 = val(args, val_loader, model)
    dice_avg, dice_1, dice_2, dice_3, dice_list, auc, auc_1, auc_2, auc_3 = val(args, val_loader, model)

    return dice_avg, dice_1, dice_2, dice_3, dice_list, auc, auc_1, auc_2, auc_3


def parse_args():
    parser = argparse.ArgumentParser(description='train')
    # config
    parser.add_argument('-d', '--data-dir', default='./data/dataset/')
    parser.add_argument('-l', '--list-dir', default='./data/data_path/',
                        help='List dir to look for train_images.txt etc. '
                             'It is the same with --data-dir if not set.')
    parser.add_argument('--name', dest='name', help='change model', default='unet', type=str)
    parser.add_argument('--resume', default='', type=str, metavar='PATH',
                        help='path to latest checkpoint (default: none)')
    parser.add_argument('-j', '--workers', type=int, default=0)
    # Train Setting
    parser.add_argument('--step', type=int, default=200)
    parser.add_argument('--batch-size', type=int, default=8, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--epochs', type=int, default=30, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--lr-mode', type=str, default='step')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                        help='SGD momentum (default: 0.9)')
    parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                        metavar='W', help='weight decay (default: 1e-4)')
    parser.add_argument('--loss', help='change model', default='ce=[0.01,20,20,20]+cls_bce+inter',#+seg_con_loss_L1',#+con_y+intra',
                        type=str)  # ce=[0.01,20,20]#ce[0.0003434,20,20]
    parser.add_argument('-o', '--optimizer', default='Adam', type=str)
    # Data Transform
    parser.add_argument('--random-rotate', default=0, type=int)
    parser.add_argument('--random-scale', default=0, type=float)
    parser.add_argument('--resize', default=[512,496], type=int)
    parser.add_argument('-s', '--crop-size', default=0, type=int)
    parser.add_argument('--supervision', default='point', help='full or point')
    # Pretrain and Checkpoint
    parser.add_argument('-p', '--pretrained', type=bool)
    parser.add_argument('--model-path', default=None, type=str)
    parser.add_argument('--save-every-checkpoint', action='store_true')
    parser.add_argument('--describe', default='12.28v1',
                        help='additional information')  # 9.15v3_point_expanded_DSRG_th1=.96_th2=.99

    args = parser.parse_args()

    return args


def main():
    ##### config #####
    args = parse_args()
    seed = 1234
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    print('torch version:', torch.__version__)
    ##### logger setting #####
    task_name = args.list_dir.split('/')[-1]
    pretrained = 'pre' if args.pretrained else 'nopre'
    result_path = osp.join('result', task_name, 'train',
                           args.name + '_' + pretrained + '_' + args.loss + '_' + str(
                               args.lr) + '_' + args.supervision + '_' + args.describe)
    if not exists(result_path):
        os.makedirs(result_path)
    resume = True if args.resume else False
    logger = Logger(osp.join(result_path, 'dice_epoch.txt'), title='dice', resume=resume)
    # if not resume:
    logger.set_names(
        ['Epoch', 'Dice_Train', 'Dice_Val', 'AUC', 'Dice_11', 'Dice_22', 'Dice_33', 'AUC_1', 'AUC_2', 'AUC_3'])
    # train_seg
    train_seg(args, result_path, logger)


if __name__ == '__main__':
    main()