train.py

import os
import time
import gc
import cv2 as cv
import argparse
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
plt.ion()

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from torch.autograd import Variable
from torch.utils import data

from models.fcn_xu import fcn_mul
from data_loader.data_loader_18 import MR18loader_CV
from metrics import runningScore
from loss import cross_entropy2d,loss_ce_t,  weighted_loss,  dice_loss,dice_coeff,  bce2d_hed

from models.fcn_xu import fcn_xu,fcn_xu_19,fcn_nopool,fcn_xu_dilated
from models.unet import unet
from models.PAN import PAN_seg
from models.resnet import FCN_res

from models.segnet import segnet
from models.densenet import DenseNet,DenseNetSeg
from models.tiramisu import tiramisu

def adjust_learning_rate(optimizer, epoch):
    lr = args.lr * (0.1 ** (epoch // 5))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def train(args):
    os.environ["CUDA_VISIBLE_DEVICES"]=str(args.gpu_id)
    #torch.manual_seed(1337)
    print(args)
    # setup dataloader
    t_loader=MR18loader_CV(root=args.data_path,val_num=args.val_num,is_val=False,is_transform=True,is_flip=True,is_rotate=True,is_crop=True,is_histeq=True,forest=args.num_forest)
    v_loader=MR18loader_CV(root=args.data_path,val_num=args.val_num,is_val=True,is_transform=True,is_flip=False,is_rotate=False,is_crop=True,is_histeq=True,forest=args.num_forest)
    n_classes = t_loader.n_classes
    trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=1, shuffle=True)
    valloader = data.DataLoader(v_loader, batch_size=1, num_workers=1,shuffle=True)
    # setup Metrics
    running_metrics_single = runningScore(n_classes)
    running_metrics_single_test = runningScore(4)
    # setup Model
    model=fcn_mul(n_classes=n_classes)
    vgg16 = models.vgg16(pretrained=True)
    model.init_vgg16_params(vgg16)
    model.cuda()
    # setup optimizer and loss
    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.99, weight_decay=5e-4)
    loss_ce = cross_entropy2d
    #loss_ce_weight = weighted_loss
    #loss_dc = dice_loss
    #loss_hed= bce2d_hed
    # resume
    best_iou=-100.0
    if args.resume is not None:
        if os.path.isfile(args.resume):
            print("Loading model and optimizer from checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            best_iou=checkpoint['best_iou']
            model.load_state_dict(checkpoint['model_state'])
            optimizer.load_state_dict(checkpoint['optimizer_state'])
            print("Loaded checkpoint '{}' (epoch {}), best_iou={}"
                  .format(args.resume, checkpoint['epoch'],best_iou))
        else:
            best_iou=-100.0
            print("No checkpoint found at '{}'".format(args.resume))
    # visualization
    t = []
    loss_seg_list=[]
    loss_hed_list=[]
    Dice_mean=[]
    Dice_CSF=[]
    Dice_GM=[]
    Dice_WM=[]
    t_pre=time.time()
    print('training prepared, cost {} seconds\n\n'.format(t_pre-t_begin))
    for epoch in range(args.n_epoch):
        t.append(epoch+1)
        model.train()
        adjust_learning_rate(optimizer,epoch)
        #loss_sum=0.0
        loss_epoch=0.0
        t_epoch=time.time()
        for i_train, (regions,T1s,IRs,T2s,lbls) in enumerate(trainloader):
            T1s=Variable(T1s.cuda())
            IRs,T2s=Variable(IRs.cuda()),Variable(T2s.cuda())
            lbls=Variable(lbls.cuda()[:,int(args.num_forest/2),:,:].unsqueeze(1))
            #edges=Variable(edges.cuda()[:,int(args.num_forest/2),:,:].unsqueeze(1))
            optimizer.zero_grad()
            outputs=model(T1s,IRs,T2s)
            seg_out=F.log_softmax(outputs,dim=1)
            max_prob,_=torch.max(seg_out,dim=1)
            max_prob=-max_prob.detach().unsqueeze(1)
            loss_seg_value=loss_ce(input=outputs,target=lbls) 
                    #+0.5*loss_dc(input=outputs,target=lbls)
                    #+0.5*loss_ce_weight(input=outputs,target=lbls,weight=max_prob)\
                    #+0.5*loss_ce_weight(input=outputs,target=lbls,weight=edges)\
            #loss_hed_value=loss_hed(input=outputs[1],target=edges)
                #+0.5*loss_hed(input=outputs[2],target=edges) \
                #+0.5*loss_hed(input=outputs[3],target=edges) \
                #+0.5*loss_hed(input=outputs[4],target=edges) \
                #+0.5*loss_hed(input=outputs[5],target=edges)
            loss=loss_seg_value
            #loss=loss_seg_value+loss_hed_value
            # loss average
            #loss_sum+=loss
            #if (i_train+1)%args.loss_avg==0:
            #    loss_sum/=args.loss_avg
            #    loss_sum.backward()
            #    optimizer.step()
            #    loss_sum=0.0
            loss.backward()
            optimizer.step()
            loss_epoch+=loss.item()
            # visualization
            if i_train==40:
                ax1=plt.subplot(241)
                ax1.imshow(T1s[0,1,:,:].data.cpu().numpy(),cmap ='gray')
                ax1.set_title('train_img')
                ax1.axis('off')
                ax2=plt.subplot(242)
                ax2.imshow(t_loader.decode_segmap(lbls[0,0,:,:].data.cpu().numpy()).astype(np.uint8))
                ax2.set_title('train_label')
                ax2.axis('off')
                ax3=plt.subplot(243)
                model.eval()
                train_show=model(T1s,IRs,T2s)
                ax3.imshow(t_loader.decode_segmap(train_show[0].data.max(0)[1].cpu().numpy()).astype(np.uint8))
                ax3.set_title('train_predict')
                ax3.axis('off')
                ax4=plt.subplot(244)
                ax4.imshow(max_prob[0,0].cpu().numpy())
                ax4.set_title('uncertainty')
                ax4.axis('off')
                model.train()
        loss_epoch/=i_train
        loss_seg_list.append(loss_epoch)
        loss_hed_list.append(0)
        t_train=time.time()
        print('epoch: ',epoch+1)
        print('--------------------------------Training--------------------------------')
        print('average loss in this epoch: ',loss_epoch)
        print('final loss in this epoch: ',loss.data.item())
        print('cost {} seconds up to now'.format(t_train-t_begin))
        print('cost {} seconds in this train epoch'.format(t_train-t_epoch))

        model.eval()
        for i_val, (regions_val,T1s_val,IRs_val,T2s_val,lbls_val) in enumerate(valloader):
            T1s_val=Variable(T1s_val.cuda())
            IRs_val,T2s_val=Variable(IRs_val.cuda()),Variable(T2s_val.cuda())
            with torch.no_grad():
                outputs_single=model(T1s_val,IRs_val,T2s_val)[0,:,:,:]
            # get predict
            pred_single=outputs_single.data.max(0)[1].cpu().numpy()
            # pad to 240
            pred_pad=np.zeros((240,240),np.uint8)
            pred_pad[regions_val[0]:regions_val[1],regions_val[2]:regions_val[3]]=  \
                    pred_single[0:regions_val[1]-regions_val[0],0:regions_val[3]-regions_val[2]]
            # convert to 3 classes
            pred_single_test=np.zeros((240,240),np.uint8)
            pred_single_test=v_loader.lbl_totest(pred_pad)
            # get gt
            gt = lbls_val[0][int(args.num_forest/2)].numpy()
            # pad to 240
            gt_pad=np.zeros((240,240),np.uint8)
            gt_pad[regions_val[0]:regions_val[1],regions_val[2]:regions_val[3]]=  \
                    gt[0:regions_val[1]-regions_val[0],0:regions_val[3]-regions_val[2]]
            # convert to 3 classes
            gt_test=np.zeros((240,240),np.uint8)
            gt_test=v_loader.lbl_totest(gt_pad)
            # metrics update
            running_metrics_single.update(gt_pad, pred_pad)
            running_metrics_single_test.update(gt_test, pred_single_test)
            # visualization
            if i_val==40:
                ax5=plt.subplot(245)
                ax5.imshow((T1s_val[0,int(args.num_forest/2),:,:].data.cpu().numpy()*255+t_loader.T1mean).astype(np.uint8),cmap ='gray')
                ax5.set_title('src_img')
                ax5.axis('off')
                ax6=plt.subplot(246)
                ax6.imshow(t_loader.decode_segmap(gt).astype(np.uint8))
                ax6.set_title('gt')
                ax6.axis('off')
                ax7=plt.subplot(247)
                ax7.imshow(t_loader.decode_segmap(pred_single).astype(np.uint8))
                ax7.set_title('pred_single')
                ax7.axis('off')
                ax8=plt.subplot(248)
                ax8.imshow(pred_single_test[regions_val[0]:regions_val[1],regions_val[2]:regions_val[3]].astype(np.uint8))
                ax8.set_title('pred_single_test')
                ax8.axis('off')
                plt.tight_layout()
                plt.subplots_adjust(wspace=.1,hspace=.3)
                plt.savefig('./fig_out/val_{}_out_{}.png'.format(str(args.val_num),epoch+1))
        # compute dice coefficients during validation
        score_single, class_iou_single = running_metrics_single.get_scores()
        score_single_test, class_iou_single_test = running_metrics_single_test.get_scores()
        Dice_mean.append(score_single['Mean Dice : \t'])
        Dice_CSF.append(score_single_test['Dice : \t'][1])
        Dice_GM.append(score_single_test['Dice : \t'][2])
        Dice_WM.append(score_single_test['Dice : \t'][3])
        print('--------------------------------All tissues--------------------------------')
        print('Back: Background,')
        print('GM: Cortical GM(red), Basal ganglia(green),')
        print('WM: WM(yellow), WM lesions(blue),')
        print('CSF: CSF(pink), Ventricles(light blue),')
        print('Back: Cerebellum(white), Brainstem(dark red)')
        print('single predict: ')
        for k, v in score_single.items():
            print(k, v)
        print('--------------------------------Only tests--------------------------------')
        print('tissue : Back , CSF , GM , WM')
        print('single predict: ')
        for k, v in score_single_test.items():
            print(k, v)
        t_test=time.time()
        print('cost {} seconds up to now'.format(t_test-t_begin))
        print('cost {} seconds in this validation epoch'.format(t_test-t_train))
        # save model at best validation metrics
        if score_single['Mean Dice : \t'] >= best_iou:
            best_iou = score_single['Mean Dice : \t']
            state = {'epoch': epoch+1,
                     'model_state': model.state_dict(),
                     'optimizer_state' : optimizer.state_dict(),
                     'best_iou':best_iou}
            torch.save(state, "val_{}_best.pkl".format(str(args.val_num)))
            print('model saved!!!')
        # save model every 10 epochs
        if (epoch+1)%10==0:
            state = {'epoch': epoch+1,
                     'model_state': model.state_dict(),
                     'optimizer_state' : optimizer.state_dict(),
                     'score':score_single}
            torch.save(state, "val_{}_e_{}.pkl".format(str(args.val_num),epoch+1))
        # plot curve
        ax1=plt.subplot(211)
        ax1.plot(t,loss_seg_list,'g')
        ax1.plot(t,loss_hed_list,'r')
        ax1.set_title('train loss')
        ax2=plt.subplot(212)
        ax2.plot(t,Dice_mean,'k')
        ax2.plot(t,Dice_CSF,'r')
        ax2.plot(t,Dice_GM,'g')
        ax2.plot(t,Dice_WM,'b')
        ax2.set_title('validate Dice, R/G/B for CSF/GM/WM')
        plt.tight_layout()
        plt.subplots_adjust(wspace=0,hspace=.3)
        plt.savefig('./fig_out/val_{}_curve.png'.format(str(args.val_num)))
        # metric reset
        running_metrics_single.reset()
        running_metrics_single_test.reset()
        print('\n\n')

if __name__ == '__main__':
    t_begin=time.time()
    parser = argparse.ArgumentParser(description='Hyperparams')
    parser.add_argument('--gpu_id', nargs='?', type=int, default=-1,
                        help='GPU id, -1 for cpu')
    parser.add_argument('--data_path', nargs='?', type=str, default='/home/canpi/canpi/MRBrainS18/data/',
                        help='dataset path')
    parser.add_argument('--val_num', nargs='?', type=int, default=1,
                        help='which set is left for validation')
    
    parser.add_argument('--n_epoch', nargs='?', type=int, default=20,
                        help='# of the epochs')
    parser.add_argument('--batch_size', nargs='?', type=int, default=1,
                        help='Batch Size')
    parser.add_argument('--num_forest', nargs='?', type=int, default=3,
                        help='number of stacked slice')
    #parser.add_argument('--loss_avg', nargs='?', type=int, default=1,
    #                    help='loss average')
    parser.add_argument('--lr', nargs='?', type=float, default=1e-3,
                        help='Learning Rate')
    parser.add_argument('--resume', nargs='?', type=str, default=None,
                        help='Path to previous saved model to restart from')
    args = parser.parse_args()
    train(args)