miccai_main_warmup_pipsUS_echo.py

"""
Always keep the first frame as reference (regularization)
"""
import numpy as np
import saverloader
from nets.pipsUS import PipsUS
import utils.improc
import utils.geom
import utils.misc
import torch
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
from torch.utils.data import Dataset, DataLoader
from ultrasound.sanity_check_echo_pseudo_label import generate_pseudo_gt
from ultrasound.sanity_check_echodata import EchoUSDataset
import random

IMAGE_SIZE = 256
USE_BATCH = False # set to True somehow makes the model predict ~0
USE_MINI = True

def pt_sequence_loss(pt_preds, pt_gt, gamma=0.8):
    """ Loss function defined over sequence of flow predictions """
    B, N, D = pt_gt.shape
    assert(D==2)
    n_predictions = len(pt_preds)    
    flow_loss = 0.0
    # generate mask for invalid point
    mask = (pt_gt[:, :, 0] >= 0) & (pt_gt[:,:,1] >= 0) & (pt_gt[:,:,0] <IMAGE_SIZE) & (pt_gt[:,:,1] < IMAGE_SIZE) 
    mask = mask.unsqueeze(2)

    for i in range(n_predictions):
        i_weight = gamma**(n_predictions - i - 1)
        i_loss = F.l1_loss(pt_preds[i] * mask, pt_gt * mask, reduction='sum')
        flow_loss += i_weight * i_loss / torch.sum(mask)
    flow_loss = flow_loss/n_predictions
    # print("flow_loss", flow_loss)

    if torch.isnan(flow_loss):
        # output debug info
        # print("n_predictions:", n_predictions, "mask sum:", torch.sum(mask))
        # pred is nan
        if torch.sum(mask) > 0:
            for i in range(n_predictions):
                print("Iteration:", i, "pred has NaN:", torch.isnan(pt_preds[i]).any()) # HAS NaN HERE!!!!! OUTPUT EXPLODED :( https://discuss.pytorch.org/t/why-my-model-returns-nan/24329/4
    return flow_loss


def requires_grad(parameters, flag=True):
    for p in parameters:
        p.requires_grad = flag

def fetch_optimizer(lr, wdecay, epsilon, num_steps, params):
    optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=wdecay, eps=epsilon)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, lr, num_steps+100, pct_start=0.1, cycle_momentum=False, anneal_strategy='linear')
    return optimizer, scheduler


def val_model(student_model, data, device, sequence_length, iters=8, sw=None, batch_size=4):
    metrics = {}
    videos = data['rgbs'][0]
    motion = data['motion'][0]
    tracking_dataset = generate_pseudo_gt(videos, motion, is_train=False)

    if tracking_dataset.__len__() == 0:
        print("dataset length is 0! exit")
        return {'total_loss': 0}

    if USE_BATCH:
        dataloader = DataLoader(tracking_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
    else:
        dataloader = DataLoader(tracking_dataset, batch_size=1, shuffle=True, num_workers=0)
    _, _, H, W = videos.shape

    # metric to calculate
    total_loss = 0
    metrics = {}
    
    student_model.to(device)
    student_model.eval()

    with torch.no_grad():
        
        # start iteration
        for i, data in enumerate(dataloader):
            
            rgbs = data['images'] # B,video_length,C,H,W
            trajs_g = data['trajs_gt'] # B,video_length,N,2
            
            # pad with static start
            rgbs_pad = rgbs[:,0:1].repeat(1,sequence_length-1,1,1,1)
            trajs_g_pad = trajs_g[:,0:1].repeat(1,sequence_length-1,1,1)
            rgbs = torch.cat((rgbs_pad,rgbs),dim=1)
            trajs_g = torch.cat((trajs_g_pad, trajs_g),dim=1)

            B, S, C, H, W = rgbs.shape
            assert(C==3)
            B, S, N, D = trajs_g.shape
            assert(D==2)

            loss = torch.tensor(0.0).to(device)
            rgbs = rgbs.to(device)
            trajs_g = trajs_g.to(device)
            valid_loss_counter = 0
            
            for jj in range(S-sequence_length-1):

                if jj == 0:
                    trajs_previous_from_model = trajs_g[:,:sequence_length] # buffer to save the prediction from model
                    trajs_previous = trajs_g[:,:sequence_length]
                    image_previous = rgbs[:,jj:jj+sequence_length]
                else:
                    trajs_previous = torch.cat((trajs_g[:,0:1], trajs_g[:,jj+1:jj+sequence_length]), dim=1)
                    image_previous = torch.cat((rgbs[:,0:1],rgbs[:,jj+1:jj+sequence_length]), dim=1)

                preds_coords = student_model(trajs_previous, image_previous=image_previous, image_curr=rgbs[:,jj+sequence_length], iters=iters, beautify=True)
                preds_e = preds_coords[-1] # prediction at the last iteration, B,N,2
                trajs_previous_from_model = torch.cat((trajs_previous_from_model[:,1:], preds_e.detach().unsqueeze(1)), dim=1)

                # for now just MSE, in the future add regularization
                curr_loss = pt_sequence_loss(preds_coords, trajs_g[:,jj+sequence_length])

                if torch.isnan(curr_loss):
                    print('nan in loss; skipping')
                    continue
                
                loss += curr_loss
                valid_loss_counter += 1

            if valid_loss_counter > 0:
                loss = loss / valid_loss_counter
            total_loss = loss.item() + total_loss
    
        total_loss = total_loss / len(dataloader)

    # # visualize current training for the last batch
    # if sw is not None and sw.save_this:
    #     prep_rgbs = utils.improc.preprocess_color(rgbs)

    #     trajs_pred = trajs_g.clone()
    #     trajs_pred[:,-1] = preds_e[:]
    #     sw.summ_traj2ds_on_rgbs('valid/trajs_pred_on_rgbs', trajs_pred[0:1], utils.improc.preprocess_color(rgbs[0:1]), cmap='hot', linewidth=1, show_dots=False)
    #     sw.summ_traj2ds_on_rgbs('valid/trajs_gt_on_rgbs',  trajs_g[0:1], utils.improc.preprocess_color(rgbs[0:1]), cmap='hot', linewidth=1, show_dots=False)

    #     sw.summ_traj2ds_on_rgb('valid/trajs_pred_on_rgb_curr', trajs_pred[0:1], prep_rgbs[0:1,-1], cmap='spring', linewidth=2)
    #     sw.summ_traj2ds_on_rgb('valid/trajs_gt_on_rgb_curr', trajs_g[0:1], prep_rgbs[0:1,-1], cmap='spring', linewidth=2)


    # analyze stats for this run
    metrics['total_loss'] = total_loss
    student_model.to('cpu')

    return metrics

def train_model(student_model, data, device, sequence_length, optimizer, scheduler=None, iters=8, sw=None, use_augs=True, batch_size=4):

    videos = data['rgbs'][0]
    motion = data['motion'][0]

    total_loss = 0
    metrics = {}
    
    # use teacher model to get the ground truth
    # use two times of the sequence, and use schedule sampling to insert model pred into trajs_g
    tracking_dataset = generate_pseudo_gt(videos, motion, is_train=True)
    if tracking_dataset.__len__() == 0:
        print("dataset length is 0! exit")
        return 0, metrics
    
    if USE_BATCH:
        dataloader = DataLoader(tracking_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
    else:
        dataloader = DataLoader(tracking_dataset, batch_size=1, shuffle=True, num_workers=0)

    # start iteration
    student_model.to(device)
    student_model.train()

    for i, data in enumerate(dataloader):
        seq_loss = 0
        rgbs = data['images'] # B,21,C,H,W
        trajs_g = data['trajs_gt'] # B,21,N,2

        rgbs_pad = rgbs[:,0:1].repeat(1,sequence_length-1,1,1,1)
        trajs_g_pad = trajs_g[:,0:1].repeat(1,sequence_length-1,1,1)
        rgbs = torch.cat((rgbs_pad,rgbs),dim=1)
        trajs_g = torch.cat((trajs_g_pad, trajs_g),dim=1)


        if use_augs:
            if np.random.rand() < 0.5: # rot90 aug
                rgbs = rgbs.permute(0,1,2,4,3) # swap xy
                trajs_g = trajs_g.flip([3]) # swap xy

            if np.random.rand() < 0.5: # time inverse
                rgbs = rgbs.flip([1])
                trajs_g = trajs_g.flip([1])


        B, S, C, H, W = rgbs.shape
        # print("video length", S)
        assert(C==3)
        B, S, N, D = trajs_g.shape
        assert(D==2)
        # print("video length", S)

        loss = torch.tensor(0.0).to(device)
        rgbs = rgbs.to(device)
        trajs_g = trajs_g.to(device)

        valid_loss_counter = 0
        for jj in range(S-sequence_length):

            if jj == 0:
                image_previous = rgbs[:,jj:jj+sequence_length]
                trajs_previous = trajs_g[:,:sequence_length]
                preds_coords = student_model(trajs_previous, image_previous=image_previous, image_curr=rgbs[:,jj+sequence_length], iters=iters, beautify=True)
            else:
                image_previous = torch.cat((rgbs[:,0:1],rgbs[:,jj+1:jj+sequence_length]), dim=1)

                if np.random.rand() < 0.7:
                    trajs_previous_ = torch.cat((trajs_g[:,0:1], trajs_g[:,jj+1:jj+sequence_length]), dim=1)
                    # add noise?
                    trajs_previous_[:,1:] = trajs_previous_[:,1:] + torch.from_numpy(np.random.normal(0, 1, trajs_previous_[:,1:].shape)).float().to(trajs_previous_.device)
                else:
                    trajs_previous_ = torch.cat((trajs_g[:,0:1], trajs_previous[:,1:]), dim=1)
                preds_coords = student_model(trajs_previous_, image_previous=image_previous, image_curr=rgbs[:,jj+sequence_length], iters=iters, beautify=True)

            preds_e = preds_coords[-1] # prediction at the last iteration, B,N,2
            # update trajs previous
            trajs_previous = torch.cat((trajs_previous[:,1:], preds_e.detach().unsqueeze(1)), dim=1)
            loss = pt_sequence_loss(preds_coords, trajs_g[:,jj+sequence_length])

            if torch.isnan(loss):
                # print('nan in loss; skipping')
                continue

            loss.backward()
            seq_loss = seq_loss + loss.item()
            valid_loss_counter += 1
            torch.nn.utils.clip_grad_norm_(student_model.parameters(), 5.0)
            optimizer.step()
            optimizer.zero_grad()
            if scheduler is not None:
                scheduler.step()

            if device == 'cuda:0':
                torch.cuda.empty_cache()

        if valid_loss_counter > 0:
            seq_loss = seq_loss / valid_loss_counter
        total_loss += seq_loss

    total_loss = total_loss / len(dataloader)


    student_model.to('cpu')
    metrics['total_loss'] = total_loss

    # # visualize current training for the last batch
    # if sw is not None and sw.save_this:
    #     prep_rgbs = utils.improc.preprocess_color(rgbs)

    #     trajs_pred = trajs_g.clone()
    #     trajs_pred[:,-1] = preds_e[:]
    #     sw.summ_traj2ds_on_rgbs('training/trajs_pred_on_rgbs', trajs_pred[0:1], utils.improc.preprocess_color(rgbs[0:1]), cmap='hot', linewidth=1, show_dots=False)
    #     sw.summ_traj2ds_on_rgbs('training/trajs_gt_on_rgbs',  trajs_g[0:1], utils.improc.preprocess_color(rgbs[0:1]), cmap='hot', linewidth=1, show_dots=False)

    #     sw.summ_traj2ds_on_rgb('training/trajs_pred_on_rgb_curr', trajs_pred[0:1], prep_rgbs[0:1,-1], cmap='spring', linewidth=2)
    #     sw.summ_traj2ds_on_rgb('training/trajs_gt_on_rgb_curr', trajs_g[0:1], prep_rgbs[0:1,-1], cmap='spring', linewidth=2)

    return total_loss, metrics
    

def train(
        S=5, # seqlen
        stride=8, # spatial stride of the model 
        iters=6, # inference steps of the model
        use_augs=True,
        reshape_size=(IMAGE_SIZE,IMAGE_SIZE), # size of the input to the model
        keypoint = 'sift',
        # optimization
        lr=5e-4,
        use_scheduler=False,
        max_epoch=10,
        # summaries
        log_dir='./logs_train',
        log_freq=5,
        backup_freq=5,
        # saving/loading
        ckpt_dir='./checkpoints',
        keep_latest=2,
        init_dir='', # previous checkpoint to initialize with
        load_optimizer=True,
        load_step=True,
        ignore_load=None,
):
    device = 'cuda:0'

    exp_name = 'Feb27_warmup'
        
    if init_dir:
        init_dir = '%s/%s' % (ckpt_dir, init_dir)
        
    # autogen a descriptive name
    model_name = "pipsUSMICCAI_echo"
    model_name += "_i%d" % (iters)
    model_name += "_S%d" % (S)
    model_name += "_size%d_%d" % (reshape_size[0], reshape_size[1])
    model_name += "_kp%s" % (keypoint)

    lrn = "%.1e" % lr # e.g., 5.0e-04
    lrn = lrn[0] + lrn[3:5] + lrn[-1] # e.g., 5e-4
    model_name += "_lr%s" % lrn
    if use_scheduler:
        model_name += "_s"
    if use_augs:
        model_name += "_A"
    model_name += "_%s" % exp_name

    print('model_name', model_name)
    
    ckpt_path = '%s/%s' % (ckpt_dir, model_name)
    writer_t = SummaryWriter(log_dir + '/' + model_name + '/t', max_queue=10, flush_secs=60)

    # load dataset
    print("loading data...")
    dataset_t = EchoUSDataset('train', reshape_size, use_mini=USE_MINI)
    dataset_v = EchoUSDataset('val', reshape_size, use_mini=USE_MINI)
    dataloader_t = DataLoader(dataset_t, batch_size=1, shuffle=True, num_workers=0, drop_last=False)
    dataloader_v = DataLoader(dataset_v, batch_size=1, shuffle=False, num_workers=0, drop_last=False)
    print("finish loading data! Dataset size: ", len(dataset_t), "and", len(dataset_v))
    
    max_iters = max_epoch * len(dataset_t) * 40

    # setup model and optimizer
    print("setting up model and optimizer...")
    student_model = PipsUS(stride=stride) #.to(device)
    _ = saverloader.load('./reference_model', student_model)
    student_model.init_realtime_delta()
    student_model.to(device)

    parameters = list(student_model.parameters())

    weight_decay = 1e-6
    if use_scheduler:
        optimizer, scheduler = fetch_optimizer(lr, weight_decay, 1e-8, max_iters, student_model.parameters())
    else:
        optimizer = torch.optim.AdamW(parameters, lr=lr, weight_decay=weight_decay) 
        scheduler = None

    utils.misc.count_parameters(student_model)

    global_step = 0
    if init_dir:
        if load_step and load_optimizer:
            global_step = saverloader.load(init_dir, student_model, optimizer=optimizer, scheduler=scheduler, ignore_load=ignore_load)
        elif load_step:
            global_step = saverloader.load(init_dir, student_model, ignore_load=ignore_load)
        else:
            _ = saverloader.load(init_dir, student_model.module, ignore_load=ignore_load)
            global_step = 0
    
    requires_grad(parameters, True)
    student_model.train()

    best_val_l1 = 999999.999
    last_epoch = global_step
    sw_t = utils.improc.Summ_writer(
            writer=writer_t,
            global_step=global_step,
            log_freq=log_freq,
            fps=min(S,8),
            scalar_freq=log_freq//4,
            just_gif=True)

    for epoch in range(last_epoch, max_epoch):

        # training loop
        for i, data in enumerate(dataloader_t):

            if use_scheduler:
                total_loss, metrics = train_model(student_model, data, device, S, iters=iters, optimizer=optimizer, scheduler=scheduler, use_augs=use_augs, sw=None)
            else:
                total_loss, metrics = train_model(student_model, data, device, S, iters=iters, optimizer=optimizer, use_augs=use_augs, sw=None)

            if i % log_freq == 0 or i == len(dataloader_t) - 1:
                print("Training epoch ", global_step, " video ", i, "/", len(dataloader_t), ", total loss", total_loss)

            sw_t.summ_scalar('total_loss', total_loss)

            current_lr = optimizer.param_groups[0]['lr']
            sw_t.summ_scalar('_/current_lr', current_lr)
        
        global_step += 1

        saverloader.save(ckpt_path, optimizer, student_model, global_step, scheduler=scheduler, keep_latest=keep_latest)

        if global_step % backup_freq == 0:
            saverloader.save(ckpt_path, optimizer, student_model, global_step, scheduler=scheduler, keep_latest=keep_latest, model_name='backup')

        # validation loop
        val_loss = 0
        for i, data in enumerate(dataloader_v):
            student_model.eval()
            with torch.no_grad():

                metrics = val_model(student_model, data, device, S, iters=iters, sw=None)
                val_loss += metrics['total_loss']
                if i % log_freq == 0 or i == len(dataloader_v) - 1:
                    print("Valid video ", i, "/", len(dataloader_v), ", total loss", metrics['total_loss'])

        if val_loss < best_val_l1:
            saverloader.save(ckpt_path, optimizer, student_model, global_step, scheduler=scheduler, keep_latest=keep_latest, model_name='best_val')
            best_val_l1 = val_loss
            print("update best checkpoint! Current epoch: ", global_step)

            if global_step % backup_freq == 0:
                saverloader.save(ckpt_path, optimizer, student_model, global_step, scheduler=scheduler, keep_latest=keep_latest, model_name='backup_best_val')

        student_model.train()
                
                    
    writer_t.close()



if __name__ == '__main__':
    Fire(train)