}
loss_funcs = {"metric_loss": loss, "classifier_loss": classification_loss}
mining_funcs = {"tuple_miner": miner}
# We can specify loss weights if we want to. This is optional
loss_weights = {"metric_loss": 1, "classifier_loss": 0.5}
record_keeper, _, _ = logging_presets.get_record_keeper(
"example_logs", "example_tensorboard")
hooks = logging_presets.get_hook_container(record_keeper)
dataset_dict = {"val": val_dataset}
model_folder = "example_saved_models"
# Create the tester
tester = testers.GlobalEmbeddingSpaceTester(
end_of_testing_hook=hooks.end_of_testing_hook)
end_of_epoch_hook = hooks.end_of_epoch_hook(tester, dataset_dict, model_folder)
trainer = trainers.TrainWithClassifier(
models,
optimizers,
batch_size,
loss_funcs,
mining_funcs,
iterations_per_epoch,
train_dataset,
loss_weights=loss_weights,
sampler=sampler,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook)
trainer.train(num_epochs=num_epochs)
def objective(trial):
param_gen = ParameterGenerator(trial, CONF["_fix_params"], logger=logger)
# Average results of multiple folds.
print("New parameter.")
metrics = []
constructors = MODEL_DEF.get(CONF, trial, param_gen)
for i_fold, (train_dataset, dev_dataset, train_sampler,
batch_size) in enumerate(constructors["fold_generator"]()):
print(f"Fold {i_fold}")
trainer_kwargs = constructors["modules"]()
# logging
record_keeper, _, _ = logging_presets.get_record_keeper(
csv_folder=os.path.join(args.log_dir,
f"trial_{trial.number}_{i_fold}_csv"),
tensorboard_folder=os.path.join(
args.log_dir, f"trial_{trial.number}_{i_fold}_tensorboard"))
hooks = logging_presets.get_hook_container(record_keeper)
# tester
tester = testers.GlobalEmbeddingSpaceTester(
end_of_testing_hook=hooks.end_of_testing_hook,
dataloader_num_workers=args.n_test_loader)
end_of_epoch_hook = hooks.end_of_epoch_hook(
tester, {"val": dev_dataset},
os.path.join(args.log_dir, f"trial_{trial.number}_{i_fold}_model"),
test_interval=1,
patience=args.patience)
CHECKPOINT_FN = os.path.join(
args.log_dir, f"trial_{trial.number}_{i_fold}_last.pth")
def actual_end_of_epoch_hook(trainer):
continue_training = end_of_epoch_hook(trainer)
torch.save(
({k: m.state_dict()
for k, m in trainer.models.items()},
{k: m.state_dict()
for k, m in trainer.optimizers.items()},
{k: m.state_dict()
for k, m in trainer.loss_funcs.items()}, trainer.epoch),
CHECKPOINT_FN)
return continue_training
# train
if args.trainer == "MetricLossOnly":
trainer = trainers.MetricLossOnly(
batch_size=batch_size,
mining_funcs={},
dataset=train_dataset,
sampler=train_sampler,
dataloader_num_workers=args.n_train_loader,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=actual_end_of_epoch_hook,
**trainer_kwargs)
elif args.trainer == "TrainWithClassifier":
trainer = trainers.TrainWithClassifier(
batch_size=batch_size,
mining_funcs={},
dataset=train_dataset,
sampler=train_sampler,
dataloader_num_workers=args.n_train_loader,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=actual_end_of_epoch_hook,
**trainer_kwargs)
while True:
start_epoch = 1
if os.path.exists(CHECKPOINT_FN):
model_dicts, optimizer_dicts, loss_dicts, last_epoch = \
torch.load(CHECKPOINT_FN)
for k, d in model_dicts.items():
trainer.models[k].load_state_dict(d)
for k, d in optimizer_dicts.items():
trainer.optimizers[k].load_state_dict(d)
for k, d in loss_dicts.items():
trainer.loss_funcs[k].load_state_dict(d)
start_epoch = last_epoch + 1
logger.critical(f"Start from old epoch: {last_epoch + 1}")
try:
trainer.train(num_epochs=args.max_epoch,
start_epoch=start_epoch)
except Exception as err:
logger.critical(f"Error: {err}")
if not args.ignore_error:
break
else:
raise err
else:
break
rslt = hooks.get_accuracy_history(
tester, "val", metrics=["mean_average_precision_at_r"])
metrics.append(max(rslt["mean_average_precision_at_r_level0"]))
return np.mean(metrics)
if args.trainer == "MetricLossOnly":
trainer = trainers.MetricLossOnly(
batch_size=batch_size,
mining_funcs={},
dataset=train_dataset,
sampler=train_sampler,
dataloader_num_workers=32,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook,
**trainer_kwargs)
elif args.trainer == "TrainWithClassifier":
trainer = trainers.TrainWithClassifier(
batch_size=batch_size,
mining_funcs={},
dataset=train_dataset,
sampler=train_sampler,
dataloader_num_workers=32,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook,
**trainer_kwargs)
trainer.train(num_epochs=args.max_epoch)
rslt = hooks.get_accuracy_history(tester,
"val",
metrics=["mean_average_precision_at_r"])
print("Best MAP@R=", max(rslt["mean_average_precision_at_r_level0"]))
# Save best model
os.makedirs(args.model_save_dir, exist_ok=True)
model_dir = os.path.join(args.log_dir, f"model")
def train_app(cfg):
print(cfg.pretty())
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Set trunk model and replace the softmax layer with an identity function
trunk = torchvision.models.__dict__[cfg.model.model_name](pretrained=cfg.model.pretrained)
#resnet18(pretrained=True)
#trunk = models.alexnet(pretrained=True)
#trunk = models.resnet50(pretrained=True)
#trunk = models.resnet152(pretrained=True)
#trunk = models.wide_resnet50_2(pretrained=True)
#trunk = EfficientNet.from_pretrained('efficientnet-b2')
trunk_output_size = trunk.fc.in_features
trunk.fc = Identity()
trunk = torch.nn.DataParallel(trunk.to(device))
embedder = torch.nn.DataParallel(MLP([trunk_output_size, cfg.embedder.size]).to(device))
classifier = torch.nn.DataParallel(MLP([cfg.embedder.size, cfg.embedder.class_out_size])).to(device)
# Set optimizers
if cfg.optimizer.name == "sdg":
trunk_optimizer = torch.optim.SGD(trunk.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
embedder_optimizer = torch.optim.SGD(embedder.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
classifier_optimizer = torch.optim.SGD(classifier.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
elif cfg.optimizer.name == "rmsprop":
trunk_optimizer = torch.optim.RMSprop(trunk.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
embedder_optimizer = torch.optim.RMSprop(embedder.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
classifier_optimizer = torch.optim.RMSprop(classifier.parameters(), lr=cfg.optimizer.lr, momentum=cfg.optimizer.momentum, weight_decay=cfg.optimizer.weight_decay)
# Set the datasets
data_dir = os.environ["DATASET_FOLDER"]+"/"+cfg.dataset.data_dir
print("Data dir: "+data_dir)
train_dataset, val_dataset, val_samples_dataset = get_datasets(data_dir, cfg, mode=cfg.mode.type)
print("Trainset: ",len(train_dataset), "Testset: ",len(val_dataset), "Samplesset: ",len(val_samples_dataset))
# Set the loss function
if cfg.embedder_loss.name == "margin_loss":
loss = losses.MarginLoss(margin=cfg.embedder_loss.margin,nu=cfg.embedder_loss.nu,beta=cfg.embedder_loss.beta)
if cfg.embedder_loss.name == "triplet_margin":
loss = losses.TripletMarginLoss(margin=cfg.embedder_loss.margin)
if cfg.embedder_loss.name == "multi_similarity":
loss = losses.MultiSimilarityLoss(alpha=cfg.embedder_loss.alpha, beta=cfg.embedder_loss.beta, base=cfg.embedder_loss.base)
# Set the classification loss:
classification_loss = torch.nn.CrossEntropyLoss()
# Set the mining function
if cfg.miner.name == "triplet_margin":
#miner = miners.TripletMarginMiner(margin=0.2)
miner = miners.TripletMarginMiner(margin=cfg.miner.margin)
if cfg.miner.name == "multi_similarity":
miner = miners.MultiSimilarityMiner(epsilon=cfg.miner.epsilon)
#miner = miners.MultiSimilarityMiner(epsilon=0.05)
batch_size = cfg.trainer.batch_size
num_epochs = cfg.trainer.num_epochs
iterations_per_epoch = cfg.trainer.iterations_per_epoch
# Set the dataloader sampler
sampler = samplers.MPerClassSampler(train_dataset.targets, m=4, length_before_new_iter=len(train_dataset))
# Package the above stuff into dictionaries.
models = {"trunk": trunk, "embedder": embedder, "classifier": classifier}
optimizers = {"trunk_optimizer": trunk_optimizer, "embedder_optimizer": embedder_optimizer, "classifier_optimizer": classifier_optimizer}
loss_funcs = {"metric_loss": loss, "classifier_loss": classification_loss}
mining_funcs = {"tuple_miner": miner}
# We can specify loss weights if we want to. This is optional
loss_weights = {"metric_loss": cfg.loss.metric_loss, "classifier_loss": cfg.loss.classifier_loss}
schedulers = {
#"metric_loss_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(classifier_optimizer, cfg.scheduler.step_size, gamma=cfg.scheduler.gamma),
"embedder_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(embedder_optimizer, cfg.scheduler.step_size, gamma=cfg.scheduler.gamma),
"classifier_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(classifier_optimizer, cfg.scheduler.step_size, gamma=cfg.scheduler.gamma),
"trunk_scheduler_by_epoch": torch.optim.lr_scheduler.StepLR(embedder_optimizer, cfg.scheduler.step_size, gamma=cfg.scheduler.gamma),
}
experiment_name = "%s_model_%s_cl_%s_ml_%s_miner_%s_mix_ml_%02.2f_mix_cl_%02.2f_resize_%d_emb_size_%d_class_size_%d_opt_%s_lr_%02.2f_m_%02.2f_wd_%02.2f"%(cfg.dataset.name,
cfg.model.model_name,
"cross_entropy",
cfg.embedder_loss.name,
cfg.miner.name,
cfg.loss.metric_loss,
cfg.loss.classifier_loss,
cfg.transform.transform_resize,
cfg.embedder.size,
cfg.embedder.class_out_size,
cfg.optimizer.name,
cfg.optimizer.lr,
cfg.optimizer.momentum,
cfg.optimizer.weight_decay)
record_keeper, _, _ = logging_presets.get_record_keeper("logs/%s"%(experiment_name), "tensorboard/%s"%(experiment_name))
hooks = logging_presets.get_hook_container(record_keeper)
dataset_dict = {"samples": val_samples_dataset, "val": val_dataset}
model_folder = "example_saved_models/%s/"%(experiment_name)
# Create the tester
tester = OneShotTester(
end_of_testing_hook=hooks.end_of_testing_hook,
#size_of_tsne=20
)
#tester.embedding_filename=data_dir+"/embeddings_pretrained_triplet_loss_multi_similarity_miner.pkl"
tester.embedding_filename=data_dir+"/"+experiment_name+".pkl"
end_of_epoch_hook = hooks.end_of_epoch_hook(tester, dataset_dict, model_folder)
trainer = trainers.TrainWithClassifier(models,
optimizers,
batch_size,
loss_funcs,
mining_funcs,
train_dataset,
sampler=sampler,
lr_schedulers=schedulers,
dataloader_num_workers = cfg.trainer.batch_size,
loss_weights=loss_weights,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook
)
trainer.train(num_epochs=num_epochs)
tester = OneShotTester()