def get_loader(args):
transform_train = transforms.Compose([
transforms.RandomResizedCrop((args.img_size, args.img_size),
scale=(0.05, 1.0)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
transform_test = transforms.Compose([
transforms.Resize((args.img_size, args.img_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
data_path = "/lab/vislab/DATA/CUB/metric_learning"
trainset = DefaultDataset("{}/train".format(data_path), transform_train)
testset = DefaultDataset("{}/val".format(data_path), transform_test)
train_sampler = samplers.MPerClassSampler(
trainset.targets, args.m, length_before_new_iter=len(trainset))
train_loader = DataLoader(trainset,
sampler=train_sampler,
batch_size=args.train_batch_size,
num_workers=4,
pin_memory=True)
test_loader = DataLoader(testset,
shuffle=False,
batch_size=args.eval_batch_size,
num_workers=4,
pin_memory=True) if testset is not None else None
return train_loader, test_loader
def dataset_fold_generator(all=False):
folds = KFold(n_splits=conf["n_fold"], shuffle=True, random_state=1234)
dummy = np.arange(len(labels), dtype=np.int32)[:, np.newaxis]
for train_lab, dev_lab in list(
folds.split(dummy))[:conf["fold_trials"]]:
train_lab = set([labels[_] for _ in train_lab])
dev_lab = set([labels[_] for _ in dev_lab])
train_dataset = ImageFolder(root=conf["data_dir"],
transform=trans_dict["train"],
is_valid_file=lambda fn: fn.split(
os.path.sep)[-2] in train_lab)
train_sampler = samplers.MPerClassSampler(
labels=train_dataset.targets,
m=param_gen.suggest_categorical("m", conf["m_cands"]),
batch_size=conf["batch_size"],
length_before_new_iter=conf["data_per_epoch"])
dev_dataset = ImageFolder(
root=conf["data_dir"],
transform=trans_dict["valid"],
is_valid_file=lambda fn: fn.split(os.path.sep)[-2] in dev_lab)
yield train_dataset, dev_dataset, train_sampler, conf["batch_size"]
download=True)
val_dataset = datasets.CIFAR100(root="CIFAR100_Dataset",
train=False,
transform=val_transform,
download=True)
# Set the loss function
metric_loss = losses.TripletMarginLoss(margin=0.01)
synth_loss = losses.AngularLoss(alpha=35)
g_adv_loss = losses.AngularLoss(alpha=35)
# Set the mining function
miner = miners.MultiSimilarityMiner(epsilon=0.1)
# Set the dataloader sampler
sampler = samplers.MPerClassSampler(train_dataset.targets, m=4)
# Set other training parameters
batch_size = 32
num_epochs = 2
iterations_per_epoch = 100
# Set up your models, optimizers, loss functions etc.
models = {"trunk": trunk, "embedder": embedder, "generator": generator}
optimizers = {
"trunk_optimizer": trunk_optimizer,
"embedder_optimizer": embedder_optimizer,
"generator_optimizer": generator_optimizer
}
return sample
# mask = Image.open('./samples/places2/mask/mask_01.png')
transformations = transforms.Compose([
transforms.Resize((256, 256)),
# RandomMask(mask),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
dataset = datasets.ImageFolder(train_path, transformations)
train_sampler = samplers.MPerClassSampler(dataset.targets, 2, len(dataset))
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
sampler=train_sampler,
num_workers=4,
drop_last=True)
# DF-Net Utils
def resize_like(x, target, mode='bilinear'):
return F.interpolate(x, target.shape[-2:], mode=mode, align_corners=False)
def list2nparray(lst, dtype=None):
"""fast conversion from nested list to ndarray by pre-allocating space"""
def _init_state_test2(self) -> None:
"""
Initialize the state and load it from an existing checkpoint if any
"""
"""
Initialize the state and load it from an existing checkpoint if any
"""
torch.manual_seed(0)
np.random.seed(0)
print("Create data loaders", flush=True)
Input_size_Image = self._train_cfg.input_size
Test_size = Input_size_Image
print("Input size : " + str(Input_size_Image))
print("Test size : " + str(Input_size_Image))
print("Initial LR :" + str(self._train_cfg.lr))
transf = get_transforms(input_size=Input_size_Image,
test_size=Test_size,
kind='full',
crop=True,
need=('train', 'val'),
backbone=None)
transform_train = transf['train']
transform_test = transf['val']
self.train_set = datasets.ImageFolder(self._train_cfg.imnet_path +
'/train',
transform=transform_train)
self.test_set = datasets.ImageFolder(self._train_cfg.imnet_path +
'/val',
transform=transform_test)
self.train_dataset = self.train_set
self.val_dataset = self.test_set
# self.train_dataset = ClassDisjointMURA(self.train_set, transform_train)
# self.val_dataset = ClassDisjointMURA(self.test_set, transform_test)
print(
f"Total batch_size: {self._train_cfg.batch_per_gpu * self._train_cfg.num_tasks}",
flush=True)
print("Create distributed model", flush=True)
model = models.resnet152(pretrained=False)
num_ftrs = model.fc.in_features
model.fc = common_functions.Identity()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.nn.DataParallel(model.to(device))
embedder = torch.nn.DataParallel(MLP([num_ftrs, 512]).to(device))
# Set optimizers
trunk_optimizer = torch.optim.Adam(model.parameters(),
lr=0.0001,
weight_decay=0.0001)
embedder_optimizer = torch.optim.Adam(embedder.parameters(),
lr=0.0001,
weight_decay=0.0001)
# Set the loss function
loss = losses.TripletMarginLoss(margin=0.1)
# Set the mining function
miner = miners.MultiSimilarityMiner(epsilon=0.1)
# Set the dataloader sampler
self.sampler = samplers.MPerClassSampler(self.train_dataset.targets,
m=4,
length_before_new_iter=len(
self.train_dataset))
# Package the above stuff into dictionaries.
self.models_dict = {"trunk": model, "embedder": embedder}
self.optimizers = {
"trunk_optimizer": trunk_optimizer,
"embedder_optimizer": embedder_optimizer
}
self.loss_funcs = {"metric_loss": loss}
self.mining_funcs = {"tuple_miner": miner}
# Set the datasets
train_dataset = datasets.CIFAR100(root="CIFAR100_Dataset", train=True, transform=train_transform, download=True)
val_dataset = datasets.CIFAR100(root="CIFAR100_Dataset", train=False, transform=val_transform, download=True)
# Set the loss function
loss = losses.TripletMarginLoss(margin=0.01)
# Set the classification loss:
classification_loss = torch.nn.CrossEntropyLoss()
# Set the mining function
miner = miners.MultiSimilarityMiner(epsilon=0.1)
# Set the dataloader sampler
sampler = samplers.MPerClassSampler(train_dataset.targets, m=4, length_before_new_iter=3200)
# Set other training parameters
batch_size = 32
num_epochs = 2
# 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": 1, "classifier_loss": 0.5}
record_keeper, _, _ = logging_presets.get_record_keeper("example_logs", "example_tensorboard")
def train(train_data, test_data, save_model, num_epochs, lr, embedding_size,
batch_size):
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.resnet18(pretrained=True)
trunk_output_size = trunk.fc.in_features
trunk.fc = common_functions.Identity()
trunk = torch.nn.DataParallel(trunk.to(device))
# Set embedder model. This takes in the output of the trunk and outputs 64 dimensional embeddings
embedder = torch.nn.DataParallel(
MLP([trunk_output_size, embedding_size]).to(device))
# Set optimizers
trunk_optimizer = torch.optim.Adam(trunk.parameters(),
lr=lr / 10,
weight_decay=0.0001)
embedder_optimizer = torch.optim.Adam(embedder.parameters(),
lr=lr,
weight_decay=0.0001)
# Set the loss function
loss = losses.TripletMarginLoss(margin=0.1)
# Set the mining function
miner = miners.MultiSimilarityMiner(epsilon=0.1)
# Set the dataloader sampler
sampler = samplers.MPerClassSampler(train_data.targets,
m=4,
length_before_new_iter=len(train_data))
save_dir = os.path.join(
save_model, ''.join(str(lr).split('.')) + '_' + str(batch_size) + '_' +
str(embedding_size))
os.makedirs(save_dir, exist_ok=True)
# Package the above stuff into dictionaries.
models = {"trunk": trunk, "embedder": embedder}
optimizers = {
"trunk_optimizer": trunk_optimizer,
"embedder_optimizer": embedder_optimizer
}
loss_funcs = {"metric_loss": loss}
mining_funcs = {"tuple_miner": miner}
record_keeper, _, _ = logging_presets.get_record_keeper(
os.path.join(save_dir, "example_logs"),
os.path.join(save_dir, "example_tensorboard"))
hooks = logging_presets.get_hook_container(record_keeper)
dataset_dict = {"val": test_data, "train": train_data}
model_folder = "example_saved_models"
def visualizer_hook(umapper, umap_embeddings, labels, split_name, keyname,
*args):
logging.info("UMAP plot for the {} split and label set {}".format(
split_name, keyname))
label_set = np.unique(labels)
num_classes = len(label_set)
fig = plt.figure(figsize=(20, 15))
plt.title(str(split_name) + '_' + str(num_embeddings))
plt.gca().set_prop_cycle(
cycler("color", [
plt.cm.nipy_spectral(i)
for i in np.linspace(0, 0.9, num_classes)
]))
for i in range(num_classes):
idx = labels == label_set[i]
plt.plot(umap_embeddings[idx, 0],
umap_embeddings[idx, 1],
".",
markersize=1)
plt.show()
# Create the tester
tester = testers.GlobalEmbeddingSpaceTester(
end_of_testing_hook=hooks.end_of_testing_hook,
visualizer=umap.UMAP(),
visualizer_hook=visualizer_hook,
dataloader_num_workers=32,
accuracy_calculator=AccuracyCalculator(k="max_bin_count"))
end_of_epoch_hook = hooks.end_of_epoch_hook(tester,
dataset_dict,
model_folder,
test_interval=1,
patience=1)
trainer = trainers.MetricLossOnly(
models,
optimizers,
batch_size,
loss_funcs,
mining_funcs,
train_data,
sampler=sampler,
dataloader_num_workers=32,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook)
trainer.train(num_epochs=num_epochs)
if save_model is not None:
torch.save(models["trunk"].state_dict(),
os.path.join(save_dir, 'trunk.pth'))
torch.save(models["embedder"].state_dict(),
os.path.join(save_dir, 'embedder.pth'))
print('Model saved in ', save_dir)
def get_sampler(dataset):
return samplers.MPerClassSampler(dataset.targets,
m=4,
length_before_new_iter=len(dataset))
def main():
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))
# Data transformations
trans_train = transforms.Compose([
transforms.RandomApply(transforms=[
transforms.AutoAugment(transforms.AutoAugmentPolicy.IMAGENET),
# transforms.RandomPerspective(distortion_scale=0.6, p=1.0),
transforms.RandomRotation(degrees=(0, 180)),
transforms.RandomHorizontalFlip(),
]),
np.float32,
transforms.ToTensor(),
fixed_image_standardization,
])
trans_val = transforms.Compose([
# transforms.CenterCrop(120),
np.float32,
transforms.ToTensor(),
fixed_image_standardization,
])
train_dataset = datasets.ImageFolder(os.path.join(data_dir,
"train_aligned"),
transform=trans_train)
val_dataset = datasets.ImageFolder(os.path.join(data_dir, "val_aligned"),
transform=trans_val)
# Prepare the model
model = InceptionResnetV1(classify=False,
pretrained="vggface2",
dropout_prob=0.5).to(device)
# model = torch.load(model_path).to(device)
# for param in list(model.parameters())[:-8]:
# param.requires_grad = False
trunk_optimizer = torch.optim.RMSprop(model.parameters(), lr=LR)
# Set the loss function
loss = losses.TripletMarginLoss(margin=margin_p)
# Set the mining function
# miner = miners.BatchEasyHardMiner(
# pos_strategy=miners.BatchEasyHardMiner.EASY,
# neg_strategy=miners.BatchEasyHardMiner.SEMIHARD)
miner = miners.BatchHardMiner()
# 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": model}
optimizers = {"trunk_optimizer": trunk_optimizer}
loss_funcs = {"metric_loss": loss}
mining_funcs = {"tuple_miner": miner}
# Create the tester
record_keeper, _, _ = logging_presets.get_record_keeper(
"logs", "tensorboard")
hooks = logging_presets.get_hook_container(record_keeper)
dataset_dict = {"val": val_dataset, "train": train_dataset}
model_folder = "training_saved_models"
def visualizer_hook(umapper, umap_embeddings, labels, split_name, keyname,
*args):
logging.info("UMAP plot for the {} split and label set {}".format(
split_name, keyname))
label_set = np.unique(labels)
num_classes = len(label_set)
fig = plt.figure(figsize=(8, 7))
plt.gca().set_prop_cycle(
cycler("color", [
plt.cm.nipy_spectral(i)
for i in np.linspace(0, 0.9, num_classes)
]))
for i in range(num_classes):
idx = labels == label_set[i]
plt.plot(umap_embeddings[idx, 0],
umap_embeddings[idx, 1],
".",
markersize=1)
plt.show()
tester = testers.GlobalEmbeddingSpaceTester(
end_of_testing_hook=hooks.end_of_testing_hook,
dataloader_num_workers=4,
accuracy_calculator=AccuracyCalculator(
include=['mean_average_precision_at_r'], k="max_bin_count"))
end_of_epoch_hook = hooks.end_of_epoch_hook(tester,
dataset_dict,
model_folder,
patience=15,
splits_to_eval=[('val',
['train'])])
# Create the trainer
trainer = trainers.MetricLossOnly(
models,
optimizers,
batch_size,
loss_funcs,
mining_funcs,
train_dataset,
sampler=sampler,
dataloader_num_workers=8,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook)
trainer.train(num_epochs=num_epochs)
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()
