device_id = config['device_id']

if not device_id == 'cpu' and torch.cuda.is_available():

if device_id is None:

device = torch.device('cuda')

else:

if isinstance(device_id, list):

device = torch.device('cuda:{}'.format(device_id[0]))

else:

device = torch.device('cuda:{}'.format(device_id))

torch.cuda.init()

else:

device = torch.device('cpu')

# make compatible with torch dataparallel

# TODO does not work with data parallel

config['device_id'] = []

# initialize logger after observers are appended

log.initialize(_run)

cfg = _run.config

set_device(cfg)

train_cfg = cfg['training']

validation_cfg = cfg.get('validation')

checkpoint_frequency = train_cfg['checkpoint_frequency']

restore_checkpoint_cfg = train_cfg['restore_checkpoint']

max_epochs = train_cfg['epochs']

model_files = run_train(train_cfg['dataloader'], train_cfg['model'], train_cfg['scheduler'],

train_cfg['optimizer'], train_cfg['losses'], validation_cfg,

checkpoint_frequency, restore_checkpoint_cfg, max_epochs,

_run)

if 'test' in cfg:

test_dataset_cfg = cfg['test']

score = evaluate_checkpoint_on(test_dataset_cfg, model_files[-1])

log_result(score, _run)

return format_result(score)

else:

optimizer_cfg, loss_cfg, validation_cfg, checkpoint_frequency,

restore_checkpoint, max_epochs, _run):

# Lets cuDNN benchmark conv implementations and choose the fastest.

# Only good if sizes stay the same within the main loop!

torch.backends.cudnn.benchmark = True

exit_handler = ExitHandler()

device = _run.config['device']

device_id = _run.config['device_id']

# during training just one dataloader

dataloader = dataloader_builder.build(dataloader_cfg)[0]

epoch = 0

if restore_checkpoint is not None:

model_cfg, optimizer_cfg, epoch = utils.restore_checkpoint(restore_checkpoint, model_cfg, optimizer_cfg)

def overwrite(to_overwrite, dic):

to_overwrite.update(dic)

return to_overwrite

# some models depend on dataset, for example num_joints

model_cfg = overwrite(dataloader.dataset.info, model_cfg)

model = model_builder.build(model_cfg)

loss_cfg['model'] = model

loss = loss_builder.build(loss_cfg)

loss = loss.to(device)

parameters = list(model.parameters()) + list(loss.parameters())

optimizer = optimizer_builder.build(optimizer_cfg, parameters)

lr_scheduler = scheduler_builder.build(scheduler_cfg, optimizer, epoch)

if validation_cfg is None:

validation_dataloaders = None

else:

validation_dataloaders = dataloader_builder.build(validation_cfg)

keep = False

file_logger = log.get_file_logger()

logger = log.get_logger()

model = torch.nn.DataParallel(model, device_ids=device_id)

model.cuda()

model = model.train()

trained_models = []

exit_handler.register(file_logger.save_checkpoint,

model, optimizer, "atexit",

model_cfg)

start_training_time = time.time()

end = time.time()

while epoch < max_epochs:

epoch += 1

lr_scheduler.step()

logger.info("Starting Epoch %d/%d", epoch, max_epochs)

len_batch = len(dataloader)

acc_time = 0

for batch_id, data in enumerate(dataloader):

optimizer.zero_grad()

endpoints = model(data, model.module.endpoints)

logger.debug("datasets %s", list(data['split_info'].keys()))

data.update(endpoints)

# threoretically losses could also be caluclated distributed.

losses = loss(endpoints, data)

loss_mean = torch.mean(losses)

loss_mean.backward()

optimizer.step()

acc_time += time.time() - end

end = time.time()

report_after_batch(_run=_run, logger=logger, batch_id=batch_id, batch_len=len_batch,

acc_time=acc_time, loss_mean=loss_mean, max_mem=torch.cuda.max_memory_allocated())

if epoch % checkpoint_frequency == 0:

path = file_logger.save_checkpoint(model, optimizer, epoch, model_cfg)

trained_models.append(path)

report_after_epoch(_run=_run, epoch=epoch, max_epoch=max_epochs)

if validation_dataloaders is not None and \

epoch % checkpoint_frequency == 0:

model.eval()

# Lets cuDNN benchmark conv implementations and choose the fastest.

# Only good if sizes stay the same within the main loop!

# not the case for segmentation

torch.backends.cudnn.benchmark = False

score = evaluate(validation_dataloaders, model, epoch, keep=keep)

logger.info(score)

log_score(score, _run, prefix="val_", step=epoch)

torch.backends.cudnn.benchmark = True

model.train()

report_after_training(_run=_run, max_epoch=max_epochs, total_time=time.time() - start_training_time)

path = file_logger.save_checkpoint(model, optimizer, epoch, model_cfg)

if path:

trained_models.append(path)

file_logger.close()

# TODO get best performing val model

evaluate_last = _run.config['training'].get('evaluate_last', 1)

if len(trained_models) < evaluate_last:

logger.info("Only saved %d models (evaluate_last=%d)", len(trained_models), evaluate_last)

log.initialize(_run)

cfg = _run.config

set_device(cfg)

validation_cfg = cfg['validation']

restore_checkpoint_cfg = cfg['restore_checkpoint']

model_update_cfg = cfg.get('model', {})

scores = evaluate_checkpoint_on(restore_checkpoint_cfg, validation_cfg, _run, model_update_cfg)

log_score(scores, _run, "val_")

model_cfg, _, epoch = utils.restore_checkpoint(restore_checkpoint, model_cfg=model_update_cfg, map_location='cpu')

#model_cfg['backbone']['output_dim'] = 256

dataloaders = dataloader_builder.build(dataset_cfg)

model = model_builder.build(model_cfg)

# TODO needs to be from dataset

if 'seg_class_mapping' in model_cfg:

mapping = model_cfg['seg_class_mapping']

else:

mapping = None

model.seg_mapping = mapping

model = torch.nn.DataParallel(model, device_ids=_run.config['device_id'])

model = model.cuda()

import utils

import shutil

scores = {}

model = model.eval()

file_logger = log.get_file_logger()

path_prefix = os.path.join(file_logger.get_log_dir(), 'results')

for dataloader in dataloaders:

dataset = dataloader.dataset

# this is a change in design,

# dataset creates evaluation

evaluation = dataset.get_evaluation(model)

print(evaluation.name)

folder_name = os.path.join(path_prefix, str(epoch))

utils.create_dir_recursive(folder_name)

# evaluation creates writer

with evaluation.get_writer(folder_name) as writer:

for idx, data in enumerate(dataloader):

data = evaluation.before_infere(data)

endpoints = model.module.infere(data)

data_to_write = evaluation.before_saving(endpoints, data)

writer.write(**data_to_write)

print("\rDone (%d/%d)" % (idx, len(dataloader)), flush=True, end='')

score = evaluation.score()

if evaluation.name in scores:

scores[evaluation.name].update(score)

else:

scores[evaluation.name] = score

print(score)

if not keep:

logger = log.get_logger()

logger.warning("deleting evaluation files in %s", path_prefix)

shutil.rmtree(path_prefix)

import utils

model = model.eval()

file_logger = log.get_file_logger()

path_prefix = os.path.join(file_logger.get_log_dir(), 'results')

for dataloader in dataloaders:

dataset = dataloader.dataset

# this is a change in design,

# dataset creates evaluation

test_set = dataset.get_test(model)

folder_name = os.path.join(path_prefix, str(epoch))

utils.create_dir_recursive(folder_name)

# evaluation creates writer

for idx, data in enumerate(dataloader):

test_set.write(data)

pass

# logger = log.get_logger()

# cfg = _run.config

# set_device(cfg)

# train_cfg = cfg['training']

# checkpoint_frequency = train_cfg['checkpoint_frequency']

# restore_checkpoint_cfg = train_cfg['restore_checkpoint']

#

# dataloader_cfg = train_cfg['dataloader']

# model_cfg = train_cfg['model']

# scheduler_cfg = train_cfg['scheduler']

# optimizer_cfg = train_cfg['optimizer']

# loss_cfg = train_cfg['losses']

#

# device = _run.config['device']

# device_id = _run.config['device_id']

#

# dataloader = dataloader_builder.build(dataloader_cfg)

#

# model_cfg_appendix, optimizer_cfg_appendix, epoch, _ = restore_checkpoint(restore_checkpoint_cfg)

# model_cfg.update(model_cfg_appendix)

# optimizer_cfg.update(optimizer_cfg_appendix)

#

# def overwrite(to_overwrite, dic):

# to_overwrite.update(dic)

# return to_overwrite

#

# model_cfg = overwrite(dataloader.dataset.info, model_cfg)

# model = model_builder.build(model_cfg)

#

# loss = loss_builder.build(loss_cfg)

#

# parameters = list(model.parameters()) + list(loss.parameters())

# optimizer = optimizer_builder.build(optimizer_cfg, parameters)

#

#

# validation_cfg = cfg.get('validation')

# if validation_cfg is None:

# validation_dataloaders = None

# logger.warning("No validation given")

# else:

# validation_dataloaders = dataloader_builder.build(validation_cfg)

#

# if 'evaluation' in cfg:

# evaluation_cfg = cfg['evaluation']

# dataloaders, model_cfgs = evaluation_builder.build(evaluation_cfg)

# delete = evaluation_cfg['delete']

# else:

# logger.warning("No evaluation given")

#

log.initialize(_run)

new_cfg = _run.config

set_device(new_cfg)

experiment = new_cfg['experiment']

# quick fix for last saved model

num_models = new_cfg.get('last_x', 1)

model_paths = log.Logger.get_all_model_paths(experiment)

model_paths = model_paths[-num_models:]

cfg_path = log.Logger.get_cfg_path(experiment)

with open(cfg_path, 'r') as f:

cfg = json.load(f)

# overwrite cfg

# WARNING this means we are using a config that has not been

# filled with default values

cfg.update(new_cfg)

set_device(cfg)

# TODO

for model_path in model_paths:

score = evaluate_checkpoint_on(model_path, new_cfg['validation'], _run)

log_score(score, _run)

print(format_result(score))

from models import get_all_models

from samplers import get_all_multi_samplers, get_all_single_samplers

print("models")

print((',\n').join(get_all_models()))

print("single samplers")

print((',\n').join(get_all_single_samplers()))

print("multi samplers")

formatted = "Results:\n"

metrics = {}

for eval_name, score in result.items():

for metric, value in score.items():

metric_name = "{} @{}".format(metric, eval_name)

if metric_name in metrics:

metrics[metric_name].append(value)

else:

metrics[metric_name] = [value]

for metric_name, values in metrics.items():

try:

formatted += "{}: {}\n".format(metric_name, ','.join(list(values)))

except:

formatted += "{}: {}\n".format(metric_name, str(values))

for model, scores in results.items():

for eval_name, score in scores.items():

for metric, value in score.items():

metric_name = "{} @{}".format(metric, eval_name)

for eval_name, score in scores.items():

for metric, value in score.items():

metric_name = prefix + "{} @{}".format(metric, eval_name)

ex = Experiment(name)

print(name)

ex.capture(set_device)

ex.main(main)

ex.capture(run_train)

ex.command(evaluate_experiment)

ex.command(test_config)

ex.command(show_options)

ex.command(evaluate_checkpoint)

# set default values

ex = config_builder.build(ex)

# set observers but check if maybe sacred will create them

# on its own

""" TODO

Problem is that we create the experiment before the command line is parsed by sacred.

But then we cannot set default values without using a shell script. Or

modify the file path for the logger."""

if Config.LOG_DIR is not None and '-F' not in sacred_args:

path = Config.LOG_DIR

if name is not None:

path = os.path.join(path, name)

file_ob = FileStorageObserver.create(path)

ex.observers.append(file_ob)

if Config.SLACK_WEBHOOK_URL != "":

slack_ob = SlackObserver(Config.SLACK_WEBHOOK_URL)

ex.observers.append(slack_ob)

if (Config.MONGO_DB_NAME is not None and Config.MONGO_DB_NAME != "") \

and '-m' not in sacred_args:

if Config.MONGO_USER != "":

mongo_ob = MongoObserver.create(username=Config.MONGO_USER, password=Config.MONGO_PW,

url=Config.MONGO_URL, authMechanism="SCRAM-SHA-256", db_name=Config.MONGO_DB_NAME)

else:

mongo_ob = MongoObserver.create(url=Config.MONGO_URL, db_name=Config.MONGO_DB_NAME)

ex.observers.append(mongo_ob)