def run(output_path, config):
device = "cuda"
local_rank = config['local_rank']
distributed = backend is not None
if distributed:
torch.cuda.set_device(local_rank)
device = "cuda"
rank = dist.get_rank() if distributed else 0
# Rescale batch_size and num_workers
ngpus_per_node = torch.cuda.device_count()
ngpus = dist.get_world_size() if distributed else 1
batch_size = config['batch_size'] // ngpus
num_workers = int(
(config['num_workers'] + ngpus_per_node - 1) / ngpus_per_node)
train_labelled_loader, test_loader = \
get_train_test_loaders(path=config['data_path'],
batch_size=batch_size,
distributed=distributed,
num_workers=num_workers)
model = get_model(config['model'])
model = model.to(device)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[
local_rank,
], output_device=local_rank)
optimizer = optim.SGD(model.parameters(),
lr=config['learning_rate'],
momentum=config['momentum'],
weight_decay=config['weight_decay'],
nesterov=True)
criterion = nn.CrossEntropyLoss().to(device)
le = len(train_labelled_loader)
milestones_values = [(0, 0.0),
(le * config['num_warmup_epochs'],
config['learning_rate']),
(le * config['num_epochs'], 0.0)]
lr_scheduler = PiecewiseLinear(optimizer,
param_name="lr",
milestones_values=milestones_values)
def _prepare_batch(batch, device, non_blocking):
x, y = batch
return (convert_tensor(x, device=device, non_blocking=non_blocking),
convert_tensor(y, device=device, non_blocking=non_blocking))
def process_function(engine, labelled_batch):
x, y = _prepare_batch(labelled_batch, device=device, non_blocking=True)
model.train()
# Supervised part
y_pred = model(x)
loss = criterion(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return {
'batch loss': loss.item(),
}
trainer = Engine(process_function)
if not hasattr(lr_scheduler, "step"):
trainer.add_event_handler(Events.ITERATION_STARTED, lr_scheduler)
else:
trainer.add_event_handler(Events.ITERATION_COMPLETED,
lambda engine: lr_scheduler.step())
metric_names = [
'batch loss',
]
def output_transform(x, name):
return x[name]
for n in metric_names:
# We compute running average values on the output (batch loss) across all devices
RunningAverage(output_transform=partial(output_transform, name=n),
epoch_bound=False,
device=device).attach(trainer, n)
if rank == 0:
checkpoint_handler = ModelCheckpoint(dirname=output_path,
filename_prefix="checkpoint")
trainer.add_event_handler(Events.ITERATION_COMPLETED(every=1000),
checkpoint_handler, {
'model': model,
'optimizer': optimizer
})
ProgressBar(persist=True,
bar_format="").attach(trainer,
event_name=Events.EPOCH_STARTED,
closing_event_name=Events.COMPLETED)
if config['display_iters']:
ProgressBar(persist=False,
bar_format="").attach(trainer,
metric_names=metric_names)
tb_logger = TensorboardLogger(log_dir=output_path)
tb_logger.attach(trainer,
log_handler=tbOutputHandler(
tag="train", metric_names=metric_names),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=tbOptimizerParamsHandler(optimizer,
param_name="lr"),
event_name=Events.ITERATION_STARTED)
metrics = {
"accuracy": Accuracy(device=device if distributed else None),
"loss": Loss(criterion, device=device if distributed else None)
}
evaluator = create_supervised_evaluator(model,
metrics=metrics,
device=device,
non_blocking=True)
train_evaluator = create_supervised_evaluator(model,
metrics=metrics,
device=device,
non_blocking=True)
def run_validation(engine):
torch.cuda.synchronize()
train_evaluator.run(train_labelled_loader)
evaluator.run(test_loader)
trainer.add_event_handler(Events.EPOCH_STARTED(every=3), run_validation)
trainer.add_event_handler(Events.COMPLETED, run_validation)
if rank == 0:
if config['display_iters']:
ProgressBar(persist=False,
desc="Train evaluation").attach(train_evaluator)
ProgressBar(persist=False,
desc="Test evaluation").attach(evaluator)
tb_logger.attach(train_evaluator,
log_handler=tbOutputHandler(tag="train",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.COMPLETED)
tb_logger.attach(evaluator,
log_handler=tbOutputHandler(tag="test",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.COMPLETED)
# Store the best model
def default_score_fn(engine):
score = engine.state.metrics['accuracy']
return score
score_function = default_score_fn if not hasattr(
config, "score_function") else config.score_function
best_model_handler = ModelCheckpoint(
dirname=output_path,
filename_prefix="best",
n_saved=3,
global_step_transform=global_step_from_engine(trainer),
score_name="val_accuracy",
score_function=score_function)
evaluator.add_event_handler(Events.COMPLETED, best_model_handler, {
'model': model,
})
trainer.run(train_labelled_loader, max_epochs=config['num_epochs'])
if rank == 0:
tb_logger.close()
def run(train_config, logger, **kwargs):
logger = logging.getLogger('UDA')
if getattr(train_config, 'debug', False):
setup_logger(logger, logging.DEBUG)
# Set Polyaxon environment if needed
plx_logger = None
save_dir = None
output_experiment_path = None
try:
plx_logger = PolyaxonLogger()
experiment = plx_logger.experiment
save_dir = get_outputs_path()
output_experiment_path = get_outputs_refs_paths()
output_experiment_path = output_experiment_path['experiments'][
0] if output_experiment_path else None
logger.debug("Experiment info: {}".format(
experiment.get_experiment_info()))
except PolyaxonClientException as e:
logger.warning('Logger Polyaxon : ' + str(e))
# Path configuration
saves_dict = getattr(train_config, 'saves', {})
save_dir = saves_dict.get('save_dir', '') if save_dir is None else save_dir
log_dir = os.path.join(save_dir, saves_dict.get('log_dir', ''))
save_model_dir = os.path.join(save_dir, saves_dict.get('model_dir', ''))
save_prediction_dir = os.path.join(save_dir,
saves_dict.get('prediction_dir', ''))
save_config_dir = os.path.join(save_dir, saves_dict.get('config_dir', ''))
load_model_file = saves_dict.get('load_model_file', '')
load_optimizer_file = saves_dict.get('load_optimizer_file', '')
# Create folders
create_save_folders(save_dir, saves_dict)
if output_experiment_path is not None:
model_dir = saves_dict.get('model_dir', '')
load_model_file = os.path.join(
output_experiment_path, model_dir,
load_model_file) if load_model_file else None
load_optimizer_file = os.path.join(
output_experiment_path, model_dir,
load_optimizer_file) if load_optimizer_file else None
num_epochs = getattr(train_config, 'num_epochs')
num_classes = getattr(train_config, 'num_classes')
device = getattr(train_config, 'device', 'cpu')
# Set magical acceleration
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
else:
assert device == 'cpu', 'CUDA device selected but none is available'
# Set half precision if required
use_fp_16 = getattr(train_config, 'use_fp_16', False)
train1_sup_loader = getattr(train_config, 'train1_sup_loader')
train1_unsup_loader = getattr(train_config, 'train1_unsup_loader')
train2_unsup_loader = getattr(train_config, 'train2_unsup_loader')
test_loader = getattr(train_config, 'test_loader')
save_interval = saves_dict.get('save_interval', 0)
n_saved = saves_dict.get('n_saved', 0)
val_interval = getattr(train_config, 'val_interval', 1)
pred_interval = getattr(train_config, 'pred_interval', 0)
model = getattr(train_config, 'model').to(device)
optimizer = getattr(train_config, 'optimizer')
criterion = getattr(train_config, 'criterion').to(device)
consistency_criterion = getattr(train_config,
'consistency_criterion').to(device)
cm_metric = getattr(
train_config, 'cm_metric',
ConfusionMatrix(num_classes=num_classes,
output_transform=lambda x: (x['y_pred'], x['y'])))
# AMP initialization for half precision
if use_fp_16:
assert 'cuda' in device
assert torch.backends.cudnn.enabled, "NVIDIA/Apex:Amp requires cudnn backend to be enabled."
try:
from apex import amp
except:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to run this example."
)
# Initialize amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
# Load checkpoint
load_params(model,
optimizer=optimizer,
model_file=load_model_file,
optimizer_file=load_optimizer_file,
device_name=device)
# Add batch norm
is_bn = getattr(train_config, 'is_bn', False)
if is_bn:
batch_norm = nn.BatchNorm2d(3).to(device)
if use_fp_16:
batch_norm = amp.initialize(batch_norm)
batch_norm.reset_parameters()
model = nn.Sequential(batch_norm, model)
# Copy the config file
shutil.copy2(os.path.abspath(train_config.__file__),
os.path.join(save_config_dir, 'checkpoint_module.py'))
le = len(train1_sup_loader)
num_train_steps = le * num_epochs
mlflow.log_param("num train steps", num_train_steps)
lr = getattr(train_config, 'learning_rate')
num_warmup_steps = getattr(train_config, 'num_warmup_steps', 0)
lr_scheduler = getattr(train_config, 'lr_scheduler', None)
if lr_scheduler is not None:
lr_scheduler = lr_scheduler(optimizer)
if num_warmup_steps > 0:
lr_scheduler = create_lr_scheduler_with_warmup(
lr_scheduler,
warmup_start_value=0.0,
warmup_end_value=lr * (1.0 + 1.0 / num_warmup_steps),
warmup_duration=num_warmup_steps)
train1_sup_loader_iter = cycle(train1_sup_loader)
train1_unsup_loader_iter = cycle(train1_unsup_loader)
train2_unsup_loader_iter = cycle(train2_unsup_loader)
# Reduce on plateau
reduce_on_plateau = getattr(train_config, 'reduce_on_plateau', None)
# Output transform model
output_transform_model = getattr(train_config, 'output_transform_model',
lambda x: x)
inference_fn = getattr(train_config, 'inference_fn', inference_standard)
lam = getattr(train_config, 'consistency_lambda')
beta = getattr(train_config, 'consistency_beta', lam)
tsa = TrainingSignalAnnealing(
num_steps=num_train_steps,
min_threshold=getattr(train_config, 'TSA_proba_min'),
max_threshold=getattr(train_config, 'TSA_proba_max'))
with_tsa = getattr(train_config, 'with_TSA', False)
cfg = {
'tsa': tsa,
'lambda': lam,
'beta': beta,
'with_tsa': with_tsa,
'device': device,
'consistency_criterion': consistency_criterion,
'criterion': criterion
}
trainer = Engine(
partial(train_update_function,
model=model,
optimizer=optimizer,
cfg=cfg,
train1_sup_loader_iter=train1_sup_loader_iter,
train1_unsup_loader_iter=train1_unsup_loader_iter,
train2_unsup_loader_iter=train2_unsup_loader_iter,
output_transform_model=output_transform_model,
use_fp_16=use_fp_16))
# Register events
for e in CustomEvents:
State.event_to_attr[e] = 'iteration'
trainer.register_events(*CustomEvents)
if with_tsa:
trainer.add_event_handler(Events.ITERATION_COMPLETED, log_tsa, tsa)
if lr_scheduler is not None:
if not hasattr(lr_scheduler, "step"):
trainer.add_event_handler(Events.ITERATION_STARTED, lr_scheduler)
else:
trainer.add_event_handler(Events.ITERATION_STARTED,
lambda engine: lr_scheduler.step())
trainer.add_event_handler(Events.ITERATION_COMPLETED, log_learning_rate,
optimizer)
metric_names = [
'supervised batch loss', 'consistency batch loss', 'final batch loss'
]
def output_transform(x, name):
return x[name]
for n in metric_names:
RunningAverage(
output_transform=partial(output_transform, name=n)).attach(
trainer, n)
ProgressBar(persist=True,
bar_format="").attach(trainer,
event_name=Events.EPOCH_STARTED,
closing_event_name=Events.COMPLETED)
# Handlers for Tensorboard logging
tb_logger = TensorboardLogger(log_dir=log_dir)
tb_logger.attach(trainer,
log_handler=tbOutputHandler(tag="train",
metric_names=metric_names),
event_name=CustomEvents.ITERATION_K_COMPLETED)
tb_logger.attach(trainer,
log_handler=tbOptimizerParamsHandler(optimizer,
param_name="lr"),
event_name=CustomEvents.ITERATION_K_STARTED)
# Handlers for Polyaxon logging
if plx_logger is not None:
plx_logger.attach(trainer,
log_handler=plxOutputHandler(
tag="train", metric_names=metric_names),
event_name=CustomEvents.ITERATION_K_COMPLETED)
metrics = {
'loss': Loss(criterion,
output_transform=lambda x: (x['y_pred'], x['y'])),
'mAcc': cmAccuracy(cm_metric).mean(),
'mPr': cmPrecision(cm_metric).mean(),
'mRe': cmRecall(cm_metric).mean(),
'mIoU': mIoU(cm_metric),
'mF1': cmFbeta(cm_metric, 1).mean()
}
iou = IoU(cm_metric)
for i in range(num_classes):
key_name = 'IoU_{}'.format(str(i))
metrics[key_name] = iou[i]
inference_update_fn = partial(
inference_update_function,
model=model,
cfg=cfg,
output_transform_model=output_transform_model,
inference_fn=inference_fn)
evaluator = Engine(inference_update_fn)
train_evaluator = Engine(inference_update_fn)
for name, metric in metrics.items():
metric.attach(train_evaluator, name)
metric.attach(evaluator, name)
# Add checkpoint
if save_model_dir:
checkpoint = ModelCheckpoint(dirname=save_model_dir,
filename_prefix='checkpoint',
save_interval=save_interval,
n_saved=n_saved,
create_dir=True)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint, {
'mymodel': model,
'optimizer': optimizer
})
def trigger_k_iteration_started(engine, k):
if engine.state.iteration % k == 0:
engine.fire_event(CustomEvents.ITERATION_K_STARTED)
def trigger_k_iteration_completed(engine, k):
if engine.state.iteration % k == 0:
engine.fire_event(CustomEvents.ITERATION_K_COMPLETED)
def run_validation(engine, validation_interval):
if (trainer.state.epoch - 1) % validation_interval == 0:
train_evaluator.run(train1_sup_loader)
evaluator.run(test_loader)
if save_prediction_dir:
train_output = train_evaluator.state.output
test_output = evaluator.state.output
iteration = str(trainer.state.iteration)
epoch = str(trainer.state.epoch)
save_prediction('train_{}_{}'.format(iteration, epoch),
save_prediction_dir,
train_output['x'],
torch.argmax(
train_output['y_pred'][0, :, :, :], dim=0),
y=train_output['y'][0, :, :])
save_prediction('test_{}_{}'.format(iteration, epoch),
save_prediction_dir,
test_output['x'],
torch.argmax(test_output['y_pred'][0, :, :, :],
dim=0),
y=test_output['y'][0, :, :])
train_evaluator.state.output = None
evaluator.state.output = None
if reduce_on_plateau is not None:
reduce_on_plateau.step(evaluator.state.metrics['mIoU'])
trainer.add_event_handler(Events.ITERATION_STARTED,
trigger_k_iteration_started,
k=10)
trainer.add_event_handler(Events.ITERATION_COMPLETED,
trigger_k_iteration_completed,
k=10)
trainer.add_event_handler(Events.EPOCH_STARTED,
run_validation,
validation_interval=val_interval)
trainer.add_event_handler(Events.COMPLETED,
run_validation,
validation_interval=1)
def trainer_prediction_save(engine, prediction_interval):
if (engine.state.iteration - 1) % prediction_interval == 0:
if save_prediction_dir:
trainer_output = trainer.state.output['unsup pred']
iteration = str(trainer.state.iteration)
epoch = str(trainer.state.epoch)
save_prediction('trainer_{}_{}'.format(iteration, epoch),
save_prediction_dir, trainer_output['x'],
trainer_output['y_pred'])
logger.debug(
'Saved trainer prediction for iteration {}'.format(
str(engine.state.iteration)))
trainer.state.output = None
trainer.add_event_handler(Events.ITERATION_COMPLETED,
trainer_prediction_save,
prediction_interval=pred_interval)
tb_logger.attach(train_evaluator,
log_handler=tbOutputHandler(tag="train",
metric_names=list(
metrics.keys())),
event_name=Events.EPOCH_COMPLETED)
tb_logger.attach(evaluator,
log_handler=tbOutputHandler(tag="test",
metric_names=list(
metrics.keys())),
event_name=Events.EPOCH_COMPLETED)
# Handlers for Polyaxon logging
if plx_logger is not None:
plx_logger.attach(train_evaluator,
log_handler=plxOutputHandler(tag="train",
metric_names=list(
metrics.keys())),
event_name=Events.EPOCH_COMPLETED)
plx_logger.attach(evaluator,
log_handler=plxOutputHandler(tag="test",
metric_names=list(
metrics.keys())),
event_name=Events.EPOCH_COMPLETED)
trainer.add_event_handler(Events.ITERATION_COMPLETED,
mlflow_batch_metrics_logging, "train", trainer)
train_evaluator.add_event_handler(Events.COMPLETED,
mlflow_val_metrics_logging, "train",
trainer)
evaluator.add_event_handler(Events.COMPLETED, mlflow_val_metrics_logging,
"test", trainer)
data_steps = list(range(len(train1_sup_loader)))
logger.debug('Start training')
trainer.run(data_steps, max_epochs=num_epochs)
logger.debug('Finished training')
def run(output_path, config):
device = "cuda"
batch_size = config['batch_size']
train_loader, test_loader = get_train_test_loaders(
dataset_name=config['dataset'],
path=config['data_path'],
batch_size=batch_size,
num_workers=config['num_workers'])
model = get_model(config['model'])
model = model.to(device)
optim_fn = optim.SGD
if config['with_layca']:
optim_fn = LaycaSGD
optimizer = optim_fn(model.parameters(),
lr=0.0,
momentum=config['momentum'],
weight_decay=config['weight_decay'],
nesterov=True)
criterion = nn.CrossEntropyLoss()
le = len(train_loader)
milestones_values = [(le * m, v)
for m, v in config['lr_milestones_values']]
scheduler = PiecewiseLinear(optimizer,
"lr",
milestones_values=milestones_values)
def _prepare_batch(batch, device, non_blocking):
x, y = batch
return (convert_tensor(x, device=device, non_blocking=non_blocking),
convert_tensor(y, device=device, non_blocking=non_blocking))
def process_function(engine, batch):
x, y = _prepare_batch(batch, device=device, non_blocking=True)
model.train()
y_pred = model(x)
loss = criterion(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return loss.item()
trainer = Engine(process_function)
trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
RunningAverage(output_transform=lambda x: x,
epoch_bound=False).attach(trainer, 'batchloss')
ProgressBar(persist=True,
bar_format="").attach(trainer,
event_name=Events.EPOCH_STARTED,
closing_event_name=Events.COMPLETED)
tb_logger = TensorboardLogger(log_dir=output_path)
tb_logger.attach(trainer,
log_handler=tbOutputHandler(tag="train",
metric_names='all'),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=tbOptimizerParamsHandler(optimizer,
param_name="lr"),
event_name=Events.ITERATION_STARTED)
tb_logger.attach(trainer,
log_handler=LayerRotationStatsHandler(model),
event_name=Events.EPOCH_STARTED)
metrics = {
"accuracy": Accuracy(),
}
evaluator = create_supervised_evaluator(model,
metrics=metrics,
device=device,
non_blocking=True)
train_evaluator = create_supervised_evaluator(model,
metrics=metrics,
device=device,
non_blocking=True)
def run_validation(engine, val_interval):
if (engine.state.epoch - 1) % val_interval == 0:
train_evaluator.run(train_loader)
evaluator.run(test_loader)
trainer.add_event_handler(Events.EPOCH_COMPLETED,
run_validation,
val_interval=2)
trainer.add_event_handler(Events.COMPLETED, run_validation, val_interval=1)
tb_logger.attach(train_evaluator,
log_handler=tbOutputHandler(tag="train",
metric_names='all',
another_engine=trainer),
event_name=Events.COMPLETED)
tb_logger.attach(evaluator,
log_handler=tbOutputHandler(tag="test",
metric_names='all',
another_engine=trainer),
event_name=Events.COMPLETED)
def mlflow_batch_metrics_logging(engine, tag):
step = trainer.state.iteration
for name, value in engine.state.metrics.items():
mlflow.log_metric("{} {}".format(tag, name), value, step=step)
def mlflow_val_metrics_logging(engine, tag):
step = trainer.state.epoch
for name in metrics.keys():
value = engine.state.metrics[name]
mlflow.log_metric("{} {}".format(tag, name), value, step=step)
trainer.add_event_handler(Events.ITERATION_COMPLETED,
mlflow_batch_metrics_logging, "train")
train_evaluator.add_event_handler(Events.COMPLETED,
mlflow_val_metrics_logging, "train")
evaluator.add_event_handler(Events.COMPLETED, mlflow_val_metrics_logging,
"test")
trainer.run(train_loader, max_epochs=config['num_epochs'])
tb_logger.close()
def run(output_path, config):
device = "cuda"
batch_size = config['batch_size']
train_labelled_loader, train_unlabelled_loader, test_loader = \
get_train_test_loaders(dataset_name=config['dataset'],
num_labelled_samples=config['num_labelled_samples'],
path=config['data_path'],
batch_size=batch_size,
unlabelled_batch_size=config.get('unlabelled_batch_size', None),
num_workers=config['num_workers'])
model = get_model(config['model'])
model = model.to(device)
optimizer = optim.SGD(model.parameters(),
lr=config['learning_rate'],
momentum=config['momentum'],
weight_decay=config['weight_decay'],
nesterov=True)
with_SWA = config['with_SWA']
if with_SWA:
optimizer = torchcontrib.optim.SWA(optimizer)
criterion = nn.CrossEntropyLoss().to(device)
if config['consistency_criterion'] == "MSE":
consistency_criterion = nn.MSELoss()
elif config['consistency_criterion'] == "KL":
consistency_criterion = nn.KLDivLoss(reduction='batchmean')
else:
raise RuntimeError("Unknown consistency criterion {}".format(
config['consistency_criterion']))
consistency_criterion = consistency_criterion.to(device)
le = len(train_labelled_loader)
num_train_steps = le * config['num_epochs']
mlflow.log_param("num train steps", num_train_steps)
lr = config['learning_rate']
eta_min = lr * config['min_lr_ratio']
num_warmup_steps = config['num_warmup_steps']
lr_scheduler = CosineAnnealingLR(optimizer,
eta_min=eta_min,
T_max=num_train_steps - num_warmup_steps)
if num_warmup_steps > 0:
lr_scheduler = create_lr_scheduler_with_warmup(
lr_scheduler,
warmup_start_value=0.0,
warmup_end_value=lr * (1.0 + 1.0 / num_warmup_steps),
warmup_duration=num_warmup_steps)
def _prepare_batch(batch, device, non_blocking):
x, y = batch
return (convert_tensor(x, device=device, non_blocking=non_blocking),
convert_tensor(y, device=device, non_blocking=non_blocking))
def cycle(iterable):
while True:
for i in iterable:
yield i
train_unlabelled_loader_iter = cycle(train_unlabelled_loader)
lam = config['consistency_lambda']
tsa = TrainingSignalAnnealing(num_steps=num_train_steps,
min_threshold=config['TSA_proba_min'],
max_threshold=config['TSA_proba_max'])
with_tsa = config['with_TSA']
with_UDA = not config['no_UDA']
def uda_process_function(engine, labelled_batch):
x, y = _prepare_batch(labelled_batch, device=device, non_blocking=True)
if with_UDA:
unsup_x, unsup_aug_x = next(train_unlabelled_loader_iter)
unsup_x = convert_tensor(unsup_x, device=device, non_blocking=True)
unsup_aug_x = convert_tensor(unsup_aug_x,
device=device,
non_blocking=True)
model.train()
# Supervised part
y_pred = model(x)
loss = criterion(y_pred, y)
supervised_loss = loss
step = engine.state.iteration - 1
if with_tsa and with_UDA:
new_y_pred, new_y = tsa(y_pred, y, step=step)
new_loss = criterion(new_y_pred, new_y)
engine.state.tsa_log = {
"new_y_pred": new_y_pred,
"loss": loss.item(),
"tsa_loss": new_loss.item()
}
supervised_loss = new_loss
# Unsupervised part
if with_UDA:
unsup_orig_y_pred = model(unsup_x).detach()
unsup_orig_y_probas = torch.softmax(unsup_orig_y_pred, dim=-1)
unsup_aug_y_pred = model(unsup_aug_x)
unsup_aug_y_probas = torch.log_softmax(unsup_aug_y_pred, dim=-1)
consistency_loss = consistency_criterion(unsup_aug_y_probas,
unsup_orig_y_probas)
final_loss = supervised_loss
if with_UDA:
final_loss += lam * consistency_loss
optimizer.zero_grad()
final_loss.backward()
optimizer.step()
return {
'supervised batch loss': supervised_loss.item(),
'consistency batch loss':
consistency_loss.item() if with_UDA else 0.0,
'final batch loss': final_loss.item(),
}
trainer = Engine(uda_process_function)
if with_UDA and with_tsa:
@trainer.on(Events.ITERATION_COMPLETED)
def log_tsa(engine):
step = engine.state.iteration - 1
if step % 50 == 0:
mlflow.log_metric("TSA threshold",
tsa.thresholds[step].item(),
step=step)
mlflow.log_metric("TSA selection",
engine.state.tsa_log['new_y_pred'].shape[0],
step=step)
mlflow.log_metric("Original X Loss",
engine.state.tsa_log['loss'],
step=step)
mlflow.log_metric("TSA X Loss",
engine.state.tsa_log['tsa_loss'],
step=step)
if not hasattr(lr_scheduler, "step"):
trainer.add_event_handler(Events.ITERATION_STARTED, lr_scheduler)
else:
trainer.add_event_handler(Events.ITERATION_STARTED,
lambda engine: lr_scheduler.step())
@trainer.on(Events.ITERATION_STARTED)
def log_learning_rate(engine):
step = engine.state.iteration - 1
if step % 50 == 0:
lr = optimizer.param_groups[0]['lr']
mlflow.log_metric("learning rate", lr, step=step)
if with_SWA:
@trainer.on(Events.COMPLETED)
def swap_swa_sgd(engine):
optimizer.swap_swa_sgd()
optimizer.bn_update(train_labelled_loader, model)
@trainer.on(Events.EPOCH_COMPLETED)
def update_swa(engine):
if engine.state.epoch - 1 > int(num_epochs * 0.75):
optimizer.update_swa()
metric_names = [
'supervised batch loss', 'consistency batch loss', 'final batch loss'
]
def output_transform(x, name):
return x[name]
for n in metric_names:
RunningAverage(output_transform=partial(output_transform, name=n),
epoch_bound=False).attach(trainer, n)
ProgressBar(persist=True,
bar_format="").attach(trainer,
event_name=Events.EPOCH_STARTED,
closing_event_name=Events.COMPLETED)
tb_logger = TensorboardLogger(log_dir=output_path)
tb_logger.attach(trainer,
log_handler=tbOutputHandler(tag="train",
metric_names=[
'final batch loss',
'consistency batch loss',
'supervised batch loss'
]),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=tbOptimizerParamsHandler(optimizer,
param_name="lr"),
event_name=Events.ITERATION_STARTED)
metrics = {
"accuracy": Accuracy(),
}
evaluator = create_supervised_evaluator(model,
metrics=metrics,
device=device,
non_blocking=True)
train_evaluator = create_supervised_evaluator(model,
metrics=metrics,
device=device,
non_blocking=True)
def run_validation(engine, val_interval):
if (engine.state.epoch - 1) % val_interval == 0:
train_evaluator.run(train_labelled_loader)
evaluator.run(test_loader)
trainer.add_event_handler(Events.EPOCH_STARTED,
run_validation,
val_interval=2)
trainer.add_event_handler(Events.COMPLETED, run_validation, val_interval=1)
tb_logger.attach(train_evaluator,
log_handler=tbOutputHandler(tag="train",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.COMPLETED)
tb_logger.attach(evaluator,
log_handler=tbOutputHandler(tag="test",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.COMPLETED)
def mlflow_batch_metrics_logging(engine, tag):
step = trainer.state.iteration
for name, value in engine.state.metrics.items():
mlflow.log_metric("{} {}".format(tag, name), value, step=step)
def mlflow_val_metrics_logging(engine, tag):
step = trainer.state.epoch
for name in metrics.keys():
value = engine.state.metrics[name]
mlflow.log_metric("{} {}".format(tag, name), value, step=step)
trainer.add_event_handler(Events.ITERATION_COMPLETED,
mlflow_batch_metrics_logging, "train")
train_evaluator.add_event_handler(Events.COMPLETED,
mlflow_val_metrics_logging, "train")
evaluator.add_event_handler(Events.COMPLETED, mlflow_val_metrics_logging,
"test")
trainer.run(train_labelled_loader, max_epochs=config['num_epochs'])