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()