""" 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

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')

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')

videos = data['rgbs'][0]

motion = data['motion'][0]

_, _, H, W = videos.shape

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)

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 = "pipsUScorrMICCAI_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()