def __init__(self, args, data, model):

"""

Initializes an experiment builder using a named tuple (args), a data provider (data), a meta learning system

(model) and a device (e.g. gpu/cpu/n)

:param args: A namedtuple containing all experiment hyperparameters

:param data: A data provider of instance MetaLearningSystemDataLoader

:param model: A meta learning system instance

:param device: Device/s to use for the experiment

"""

self.args = args

self.device = torch.device('cuda') if args.cuda else torch.device('cpu')

# Tensorboard setup

self.use_tensorboard = (self.args.mode != 'test') and (self.args.model != 'dain')

if self.use_tensorboard:

from torch.utils.tensorboard import SummaryWriter

self.writer = SummaryWriter('logs/%s' % self.args.exp_name)

self.model = model

self.total_losses = dict()

self.state = dict()

self.state['best_val_loss'] = 0.

self.state['best_val_iter'] = 0

self.state['current_iter'] = 0

self.start_epoch = 0

self.best_PSNR = 0

self.data = data(args=args, current_iter=self.state['current_iter'])

self.state['best_epoch'] = int(self.state['best_val_iter'] / self.args.total_iter_per_epoch)

self.epoch = int(self.state['current_iter'] / self.args.total_iter_per_epoch)

self.start_time = time.time()

self.epochs_done_in_this_run = 0

if self.args.resume:

self.epoch = args.start_epoch

self.state['current_iter'] = self.epoch * self.args.total_iter_per_epoch

print(self.state['current_iter'], int(self.args.total_iter_per_epoch * self.args.max_epoch))

def build_loss_summary_string(self, summary_losses, metrics):

"""

Builds a progress bar summary string given current summary losses dictionary

:param summary_losses: Current summary statistics

:return: A summary string ready to be shown to humans.

"""

output_update = ""

for key, value in summary_losses.items():

if key != 'loss' or key.find('loss_importance_vector') >= 0:

output_update += "{}: {:.4f}, ".format(key, value)

for key, value in metrics.items():

output_update += "{}: {:.4f}, ".format(key, value.avg)

return output_update

def train_iteration(self, train_sample, sample_idx, epoch_idx, total_losses, current_iter, pbar_train, do_evaluation=False):

"""

Runs a training iteration, updates the progress bar and returns the total and current epoch train losses.

:param train_sample: A sample from the data provider

:param sample_idx: The index of the incoming sample, in relation to the current training run.

:param epoch_idx: The epoch index.

:param total_losses: The current total losses dictionary to be updated.

:param current_iter: The current training iteration in relation to the whole experiment.

:param pbar_train: The progress bar of the training.

:return: Updates total_losses, train_losses, current_iter

"""

images, metadata = train_sample

data_batch = images

losses, outputs, metrics = self.model.run_train_iter(data_batch=data_batch, epoch=epoch_idx, do_evaluation=do_evaluation)

train_output_update = self.build_loss_summary_string(losses, metrics)

pbar_train.update(1)

pbar_train.set_description("training phase {} -> {}".format(self.epoch, train_output_update))

current_iter += 1

return losses, outputs, metrics, current_iter

def evaluation_iteration(self, val_sample, total_losses, pbar_val, phase):

"""

Runs a validation iteration, updates the progress bar and returns the total and current epoch val losses.

:param val_sample: A sample from the data provider

:param total_losses: The current total losses dictionary to be updated.

:param pbar_val: The progress bar of the val stage.

:return: The updated val_losses, total_losses

"""

images, metadata = val_sample

def _eval_iter(frames):

H, W = frames[0].shape[-2:]

if H * W > 5e5 or (self.args.model == 'rrin' and H * W > 3e5): # or (self.args.model == 'dain' and H * W > 1e5):

print(H, W)

if H > W:

images_0 = [im[:, :, :H//2, :] for im in frames]

images_1 = [im[:, :, H//2:, :] for im in frames]

else:

images_0 = [im[:, :, :, :W//2] for im in frames]

images_1 = [im[:, :, :, W//2:] for im in frames]

losses_0, outputs_0, metrics_0 = _eval_iter(images_0)

losses_1, outputs_1, metrics_1 = _eval_iter(images_1)

outputs = [torch.cat([outputs_0[i], outputs_1[i]], dim=2 if H > W else 3) for i in range(len(outputs_0))]

losses = losses_0

for k, v in losses_1.items():

losses[k] = (v + losses[k]) / 2

if k == 'loss':

losses[k] = losses[k].detach()

#metrics = metrics_0

#for k, v in metrics_1.items():

# metrics[k].update(val=v.avg, n=v.count)

del losses_0, losses_1, outputs_0, outputs_1#, metrics_0, metrics_1

else:

losses, outputs, metrics = self.model.run_validation_iter(data_batch=frames)

losses['loss'] = losses['loss'].detach()

return losses, outputs, None#metrics

losses, outputs, _ = _eval_iter(images)

output = outputs[0].squeeze(0).detach()

target = images[3][0].detach().cuda()

if self.args.model == 'voxelflow':

target = (target * self.model.std + self.model.mean) / 255.0

elif self.args.model == 'superslomo':

target = self.model.revNormalize(target)

metrics = {'psnr': utils.AverageMeter(), 'ssim': utils.AverageMeter()}

psnr, ssim = utils.calc_metrics(output, target)

metrics['psnr'].update(psnr)

metrics['ssim'].update(ssim)

val_output_update = self.build_loss_summary_string(losses, metrics)

pbar_val.update(1)

pbar_val.set_description("val_phase {} -> {}".format(self.epoch, val_output_update))

return losses, outputs, metrics

def test_iteration(self, test_sample, pbar_test, phase):

"""

Runs a test iteration, updates the progress bar and returns the outputs.

:param test_sample: A sample from the data provider

:param pbar_test: The progress bar of the test stage.

:return: The output interpolation

"""

images, _ = test_sample

H, W = images[0].shape[-2:]

if H * W > 5e5:

if H > W:

images_0 = [im[:, :, :H//2, :] for im in images]

images_1 = [im[:, :, H//2:, :] for im in images]

else:

images_0 = [im[:, :, :, :W//2] for im in images]

images_1 = [im[:, :, :, W//2:] for im in images]

outputs_0 = self.model.run_test_iter(data_batch=images_0)

outputs_1 = self.model.run_test_iter(data_batch=images_1)

outputs = [torch.cat([outputs_0[i], outputs_1[i]], dim=1 if H > W else 2) for i in range(len(outputs_0))]

else:

outputs = self.model.run_test_iter(data_batch=images)

pbar_test.update(1)

# pbar_test.set_description("test_phase -> {}".format(test_output_update))

return outputs

def run_experiment(self):

"""

Runs a full training experiment with evaluations of the model on the val set at every epoch.

"""

if self.args.mode == 'test':

print('Start testing')

num_test_tasks = self.data.dataset.data_length['test']

with tqdm(total=int(num_test_tasks / self.args.test_batch_size)) as pbar_test:

for _, test_sample in enumerate(self.data.get_test_batches(total_batches=int(num_test_tasks / self.args.test_batch_size))):

outputs = self.test_iteration(test_sample=test_sample,

pbar_test=pbar_test,

phase='test')

batch_size = test_sample[0][0].shape[0]

for k in range(batch_size):

imgpath1 = test_sample[1]['imgpaths'][1][k]

imgpath2 = test_sample[1]['imgpaths'][2][k]

filename1 = imgpath1.split('/')[-1]

filename2 = imgpath2.split('/')[-1]

float_ind1 = float(filename1.split('_')[-1][:-4])

float_ind2 = float(filename2.split('_')[-1][:-4])

if float_ind2 == 0:

float_ind2 = 1.0

im_path = os.path.join(self.args.data_root, '%s_%.06f.%s'

% (filename1.split('_')[0], (float_ind1 + float_ind2) / 2, self.args.img_fmt))

utils.save_image(outputs[k], im_path)

print('Test finished.')

return

elif self.args.mode == 'val':

print('Validation only')

total_losses = dict()

val_losses = dict()

metrics_accumulator = {'psnr': utils.AverageMeter(), 'ssim': utils.AverageMeter()}

num_evaluation_tasks = self.data.dataset.data_length['val']

with tqdm(total=int(num_evaluation_tasks / self.args.val_batch_size)) as pbar_val:

for _, val_sample in enumerate(self.data.get_val_batches(total_batches=int(num_evaluation_tasks / self.args.val_batch_size))):

val_losses, outputs, metrics = self.evaluation_iteration(val_sample=val_sample,

total_losses=total_losses,

pbar_val=pbar_val,

phase='val')

batch_size = val_sample[0][0].shape[0]

for k, v in metrics.items():

metrics_accumulator[k].update(v.avg, n=v.count)

for k in range(batch_size):

paths = val_sample[1]['imgpaths'][3][k].split('/')

save_dir = os.path.join('checkpoint', self.args.exp_name, self.args.dataset, paths[-3], paths[-2])

if not os.path.exists(save_dir):

utils.makedirs(save_dir)

im_path = os.path.join(save_dir, paths[-1]) # 'im4.png' for VimeoSeptuplet

utils.save_image(outputs[0][k], im_path)

del val_losses, outputs, metrics

print("%d examples processed" % metrics_accumulator['psnr'].count)

print("PSNR: %.2f, SSIM: %.4f\n" % (metrics_accumulator['psnr'].avg, metrics_accumulator['ssim'].avg))

return

with tqdm(initial=self.state['current_iter'],

total=int(self.args.total_iter_per_epoch * self.args.max_epoch)) as pbar_train:

# training main loop

while (self.state['current_iter'] < (self.args.max_epoch * self.args.total_iter_per_epoch)):

for train_sample_idx, train_sample in enumerate(self.data.get_train_batches(total_batches=int(

self.args.total_iter_per_epoch * self.args.max_epoch) - self.state['current_iter'])):

train_losses, outputs, metrics, self.state['current_iter'] = self.train_iteration(

train_sample=train_sample,

total_losses=self.total_losses,

epoch_idx=(self.state['current_iter'] / self.args.total_iter_per_epoch),

pbar_train=pbar_train,

current_iter=self.state['current_iter'],

sample_idx=self.state['current_iter'],

do_evaluation=(self.state['current_iter'] % self.args.eval_iter == 0))

# Log to Tensorboard

if self.state['current_iter'] % self.args.log_iter == 1 and self.use_tensorboard:

utils.log_tensorboard(self.writer, train_losses, metrics['psnr'].avg, metrics['ssim'].avg, None,

self.model.optimizer.param_groups[0]['lr'], self.state['current_iter'], mode='train')

# validation

if self.state['current_iter'] % self.args.total_iter_per_epoch == 0:

total_losses = dict()

val_losses = {}

metrics_accumulator = {'psnr': utils.AverageMeter(), 'ssim': utils.AverageMeter()}

num_evaluation_tasks = self.data.dataset.data_length['val']

with tqdm(total=int(num_evaluation_tasks / self.args.val_batch_size + 0.99)) as pbar_val:

for _, val_sample in enumerate(self.data.get_val_batches(total_batches=int(

num_evaluation_tasks / self.args.val_batch_size + 0.99))):

val_loss, outputs, metrics = self.evaluation_iteration(val_sample=val_sample,

total_losses=total_losses,

pbar_val=pbar_val,

phase='val')

for k, v in metrics.items():

metrics_accumulator[k].update(v.avg, n=v.count)

for k, v in val_loss.items():

if k not in val_losses.keys():

val_losses[k] = utils.AverageMeter()

val_losses[k].update(v)

del val_loss, outputs, metrics

for k, v in val_losses.items():

val_losses[k] = v.avg

if val_losses["total"] < self.state['best_val_loss']:

print("Best validation loss", val_losses["total"])

self.state['best_val_loss'] = val_losses["total"]

self.state['best_val_psnr'] = self.state['current_iter']

self.state['best_epoch'] = int(self.state['best_val_iter'] / self.args.total_iter_per_epoch)

print("validation PSNR: %.2f, SSIM: %.4f\n" % (metrics_accumulator['psnr'].avg, metrics_accumulator['ssim'].avg))

# log to TensorBoard

if self.use_tensorboard:

utils.log_tensorboard(self.writer, val_losses, metrics_accumulator['psnr'].avg, metrics_accumulator['ssim'].avg, None,

self.model.optimizer.param_groups[0]['lr'], self.state['current_iter'], mode='val')

self.epoch += 1

PSNR = metrics_accumulator['psnr'].avg

is_best = PSNR > self.best_PSNR

self.best_PSNR = max(PSNR, self.best_PSNR)

utils.save_checkpoint({

'epoch': self.epoch,

'arch': self.args,

'state_dict': self.model.state_dict(),

# 'optimizer': optimizer.state_dict(),

'best_PSNR': self.best_PSNR

}, is_best, self.args.exp_name)

self.model.scheduler.step(val_losses['total'])

self.total_losses = dict()

self.epochs_done_in_this_run += 1

#self.evaluate_middlebury()