def get_hardest_pos_neg(samples_per_class, identities_per_batch, embeddings): l2_dist = PairwiseDistance(2).cuda() tot_faces = samples_per_class * identities_per_batch dists = [] for index, embedding in enumerate(embeddings): dists.append(list(l2_dist.forward(embedding, embeddings))) pos_mask, neg_mask = get_masks(samples_per_class, identities_per_batch) pos_mask = torch.tensor(pos_mask).cuda() neg_mask = torch.tensor(neg_mask).cuda() dists = torch.tensor(dists).cuda() pos_dists = torch.mul(dists, pos_mask) hardest_pos_dists = torch.max(pos_dists, 1) neg_dists = torch.mul(dists, neg_mask) neg_dists_nonzero = neg_dists[torch.nonzero(neg_dists, as_tuple=True)] neg_dists_nonzero_shape = (tot_faces, tot_faces - samples_per_class) neg_dists_nonzero = neg_dists_nonzero.view(neg_dists_nonzero_shape[0], neg_dists_nonzero_shape[1]) hardest_neg_dists = torch.min(neg_dists_nonzero, 1) hardest_pos_embeddings = embeddings[hardest_pos_dists[1]] hardest_neg_embeddings = embeddings[hardest_neg_dists[1]] return hardest_pos_embeddings, hardest_neg_embeddings
def main(): lfw_dataroot = args.lfw model_path = args.model_path far_target = args.far_target checkpoint = torch.load(model_path) model = Resnet18Triplet(embedding_dimension=checkpoint['embedding_dimension']) model.load_state_dict(checkpoint['model_state_dict']) flag_gpu_available = torch.cuda.is_available() if flag_gpu_available: device = torch.device("cuda") else: device = torch.device("cpu") lfw_transforms = transforms.Compose([ transforms.Resize(size=224), transforms.ToTensor(), transforms.Normalize( mean=[0.5157, 0.4062, 0.3550], std=[0.2858, 0.2515, 0.2433] ) ]) lfw_dataloader = torch.utils.data.DataLoader( dataset=LFWDataset( dir=lfw_dataroot, pairs_path='../datasets/LFW_pairs.txt', transform=lfw_transforms ), batch_size=256, num_workers=2, shuffle=False ) model.to(device) model = model.eval() with torch.no_grad(): l2_distance = PairwiseDistance(2).cuda() distances, labels = [], [] progress_bar = enumerate(tqdm(lfw_dataloader)) for batch_index, (data_a, data_b, label) in progress_bar: data_a, data_b, label = data_a.cuda(), data_b.cuda(), label.cuda() output_a, output_b = model(data_a), model(data_b) distance = l2_distance.forward(output_a, output_b) # Euclidean distance distances.append(distance.cpu().detach().numpy()) labels.append(label.cpu().detach().numpy()) labels = np.array([sublabel for label in labels for sublabel in label]) distances = np.array([subdist for distance in distances for subdist in distance]) _, _, _, _, _, _, _, tar, far = evaluate_lfw(distances=distances, labels=labels, far_target=far_target) print("TAR: {:.4f}+-{:.4f} @ FAR: {:.4f}".format(np.mean(tar), np.std(tar), np.mean(far)))
def __init__(self, num_domains, batch_size, k, device): super(AdjacencyLayer, self).__init__() self.k = k self.num_domains = num_domains self.batch_size = batch_size self.pdist = PairwiseDistance(2, keepdim=False) self.device = device
def __init__(self, model, criterion, metric_ftns, optimizer, device, num_epoch, scheduler, grad_clipping, grad_accumulation_steps, early_stopping, validation_frequency, save_period, checkpoint_dir, resume_path): self.device, device_ids = self._prepare_device(device) self.model = model.to(self.device) if len(device_ids) > 1: self.model = torch.nn.DataParallel(model, device_ids=device_ids) self.criterion = criterion self.metric_ftns = metric_ftns self.optimizer = optimizer self.num_epoch = num_epoch self.scheduler = scheduler self.grad_clipping = grad_clipping self.grad_accumulation_steps = grad_accumulation_steps self.early_stopping = early_stopping self.validation_frequency = validation_frequency self.save_period = save_period self.checkpoint_dir = checkpoint_dir self.start_epoch = 0 self.best_epoch = 0 self.best_score = 0 self.l2_dist = PairwiseDistance(2) if resume_path is not None: self._resume_checkpoint(resume_path) self.train_metrics = MetricTracker('loss') self.valid_metrics = MetricTracker( 'loss', *[m.__name__ for m in self.metric_ftns])
def findWinning(self, x): if self.Units is None: val1 = None val2 = None index1 = None index2 = None else: pdist = PairwiseDistance(p=2) if self.cuda: distance_vector = pdist(self.Units.cuda(), x.cuda()) distance_vector.to('cpu') # distance_vector = pairwise_distances(self.Units.cuda(), x.cuda()) # distance_vector.to('cpu') else: distance_vector = pdist(self.Units, x) # distance_vector = pairwise_distances(self.Units, x) if self.Units.shape[0] < 2: tuples = torch.topk(distance_vector, k=1, largest=False) val1 = tuples.values[0] index1 = tuples.indices[0] val2 = None index2 = None else: tuples = torch.topk(torch.reshape(distance_vector, (-1, )), k=2, largest=False) val1 = tuples.values[0] val2 = tuples.values[1] index1 = tuples.indices[0] index2 = tuples.indices[1] return {'val1': val1, 'index1': index1, 'val2': val2, 'index2': index2}
def compute_global_error(err): pdist = PairwiseDistance(p=2) # if cuda: # mature_neurons = mature_neurons.cuda() # data = data.cuda() # err = pairwise_distances(mature_neurons, data) # err = torch.t(torch.stack(err_list, dim=1)) mean_err = torch.mean(torch.min(err, dim=0, keepdim=True).values) return mean_err
def __init__(self, input_size: int, dropout: float = None): # L2 Distance loss supports compressed streamlines super().__init__(input_size, dropout, key='l2-regression', supports_compressed_streamlines=True, loss_description="torch's pairwise distance") # Loss will be applied on the last dimension, by default in # PairWiseDistance self.loss = PairwiseDistance()
def __init__(self, input_size: int, dropout: float = None): # Prepare the dropout, Relu, loop: super().__init__(dropout) # Layers hidden_size = ceil(input_size / 2) h1 = Linear(input_size, hidden_size) h2 = Linear(hidden_size, 3) self.layers = [h1, h2] # Loss will be applied on the last dimension, by default in # PairWiseDistance self.loss = PairwiseDistance()
def compute_cosine_matrix(self, A, B): similarity = CosineSimilarity() similarity = PairwiseDistance() print(A.size(), B.size()) num_A, _ = A.size() num_B, _ = B.size() _matrix = np.zeros((num_A, num_B)) for i in range(num_A): vA = A[i, :].unsqueeze(0) vA = vA.cuda() for j in range(num_B): vB = B[j, :].unsqueeze(0) vB = vB.cuda() value = similarity(vA, vB) # print(value) _matrix[i, j] = value.item() return _matrix
def pairwise_distances(x, y, cuda=False): """ Input: x is a Nxd matrix y is an Mxd matrix Output: dist is a NxM (Nx1 with the 0-axis sum) matrix where dist[i,j] is the square norm between x[i,:] and y[j,:] if y is not given then use 'y=x'. i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2 """ if cuda: x = x.cuda() y = y.cuda() pdist = PairwiseDistance(p=2) err_list = [] for xi in x: err_list.append(pdist(y, xi)) err = torch.t(torch.stack(err_list, dim=1)) return err
def __init__(self, p, q, metric=PairwiseDistance(p=2)): if set(p.var) != set(q.var): raise ValueError( "The two distribution variables must be the same.") if len(p.var) != 1: raise ValueError( "A given distribution must have only one variable.") super().__init__(p, q) if len(p.input_var) > 0: self.input_dist = p elif len(q.input_var) > 0: self.input_dist = q else: raise NotImplementedError() self.metric = metric
def __init__(self, module: nn.Module, model_name: str, sub_folder: str, hyperparamters: dict, optimizer, gesture_list: list, callbacks: list, train_split=CVSplit(cv=0.1, random_state=0)): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") super(SiameseEMG, self).__init__( module, criterion=nn.TripletMarginLoss, optimizer=optimizer, lr=hyperparamters['lr'], max_epochs=hyperparamters['epoch'], train_split=train_split, callbacks=callbacks, device=device, iterator_train__shuffle=True, iterator_train__num_workers=4, iterator_train__batch_size=hyperparamters['train_batch_size'], iterator_valid__shuffle=False, iterator_valid__num_workers=4, iterator_valid__batch_size=hyperparamters['valid_batch_size']) self.model_name = model_name self.hyperparamters = hyperparamters self.distance = PairwiseDistance().to(self.device) self.model_path = generate_folder('checkpoints', model_name, sub_folder=sub_folder) self.clf = KNeighborsClassifier(n_neighbors=5) self.clf_fit = False self.anchors = [] self.labels = [] print('build a new model, init parameters of {}'.format(model_name)) param_dict = self.load_pretrained("pretrained_end_params.pt") self.module.load_state_dict(param_dict)
def __init__(self): from Models.Resnet34_attention import resnet34 as resnet34_cbam model_path = os.path.join(pwd, 'Model_training_checkpoints') self.model = resnet34_cbam(pretrained=True, showlayer=False, num_classes=128) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") x = [ int(i.split('_')[4]) for i in os.listdir(model_path) if version in i ] x.sort() for i in os.listdir(model_path): if (len(x) != 0) and ('epoch_' + str(x[-1]) in i) and (version in i): model_pathi = os.path.join(model_path, i) break model_pathi = os.path.join(model_path, 'model_34_triplet_mask.pt') print(model_path) if os.path.exists(model_pathi) and (version in model_pathi): if torch.cuda.is_available(): model_state = torch.load(model_pathi) else: model_state = torch.load(model_pathi, map_location='cpu') start_epoch = model_state['epoch'] print('loaded %s' % model_pathi) else: print('模型不存在!') sys.exit(0) if torch.cuda.is_available(): self.model = self.model.cuda() self.l2_distance = PairwiseDistance(2).cuda() self.is_same = False
data_transforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) train_dataloader = DataLoader(dataset=VGGFace2Dataset( training_triplets_path=triplets_path, data_dir=data_directory, transform=data_transforms), batch_size=1, shuffle=False) net = Recog_Net() net.cuda() margin = 0.05 l2_distance = PairwiseDistance(2).cuda() optimizer_model = SGD(net.parameters(), lr=0.1) for epoch in range(0, 2): triplet_loss_sum = 0 num_valid_training_triplets = 0 progress_bar = enumerate(tqdm(train_dataloader)) for batch_idx, (batch_sample) in progress_bar: anc_image = batch_sample['anc_img'].view(-1, 3, 220, 220).cuda() pos_image = batch_sample['pos_img'].view(-1, 3, 220, 220).cuda() neg_image = batch_sample['neg_img'].view(-1, 3, 220, 220).cuda() anc_embedding = net(anc_image) pos_embedding = net(pos_image) neg_embedding = net(neg_image)
def __init__(self, num_domains, batch_size): super(DistanceLayer, self).__init__() self.num_domains = num_domains self.batch_size = batch_size self.pdist = PairwiseDistance(2)
def __init__(self, margin=1): super(ContrastiveLoss, self).__init__() self.margin = margin self.pdist = PairwiseDistance(2)
def main(): lfw_dataroot = args.lfw model_path = args.model_path far_target = args.far_target batch_size = args.batch_size flag_gpu_available = torch.cuda.is_available() if flag_gpu_available: device = torch.device("cuda") print('Using GPU') else: device = torch.device("cpu") print('Using CPU') checkpoint = torch.load(model_path, map_location=device) model = Resnet18Triplet( embedding_dimension=checkpoint['embedding_dimension']) model.load_state_dict(checkpoint['model_state_dict']) # desiredFaceWidth, height and desiredLeftEye work together to produce centered aligned faces # desiredLeftEye (x,y) says how the face moves in the heightxwidth window. bigger y moves the face down # bigger x zooms-out the face and severe values migh flip the face upside down, you wan it in range [0.2;0.36] desiredFaceHeight = 448 # set non-equal width, height so the tf_resize stretch the faces, select bigger than 224 values so we don't lose too much of the image quality desiredFaceWidth = 352 desiredLeftEye = (0.28, 0.35) tf_align = transform_align( landmark_predictor_weight_path=landmark_predictor_weights, face_detector_path=face_detector_weights_path, desiredFaceWidth=desiredFaceWidth, desiredFaceHeight=desiredFaceHeight, desiredLeftEye=desiredLeftEye) lfw_transforms = transforms.Compose([ tf_align, transforms.Resize(size=(224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.6068, 0.4517, 0.3800], std=[0.2492, 0.2173, 0.2082]) ]) lfw_dataloader = torch.utils.data.DataLoader( dataset=LFWDataset(dir=lfw_dataroot, pairs_path='../datasets/LFW_pairs.txt', transform=lfw_transforms), batch_size= batch_size, # default = 256; 160 - allows running under 2GB VRAM num_workers=2, shuffle=False) model.to(device) model = model.eval() with torch.no_grad(): l2_distance = PairwiseDistance(p=2) distances, labels = [], [] progress_bar = enumerate(tqdm(lfw_dataloader)) for batch_index, (data_a, data_b, label) in progress_bar: data_a = data_a.to(device) # data_a = data_a.cuda() data_b = data_b.to(device) # data_b = data_b.cuda() output_a, output_b = model(data_a), model(data_b) distance = l2_distance.forward(output_a, output_b) # Euclidean distance distances.append(distance.cpu().detach().numpy()) labels.append(label.cpu().detach().numpy()) labels = np.array([sublabel for label in labels for sublabel in label]) distances = np.array( [subdist for distance in distances for subdist in distance]) _, _, _, _, _, _, _, tar, far = evaluate_lfw(distances=distances, labels=labels, far_target=far_target) print("TAR: {:.4f}+-{:.4f} @ FAR: {:.4f}".format( np.mean(tar), np.std(tar), np.mean(far)))
def __init__(self, margin): super(TripletLoss, self).__init__() self.margin = margin self.pdist = PairwiseDistance(2)
def main(): dataroot = "2data" datatrainroot = "2data/training_data" datatestroot = "2data/testing_data" datavalidroot = "2data/valid_data" num_triplets_train = 1280 num_triplets_test = 1280 model_architecture = "restnet18" epochs = 3333333 resume_path = "resume" batch_size = 256 num_workers = 1 embedding_dimension = 128 pretrained = True#args.pretrained optimizer = "sgd"#args.optimizer learning_rate = 0.1#args.lr margin = 0.5#args.margin start_epoch = 0 img_size = (224, 224) data_transforms = transforms.Compose([ transforms.Resize(size=img_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize( mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5] ) ]) # Set dataloaders train_dataloader = dataloaderTriplet(datatrainroot, data_transforms, num_triplets =num_triplets_train, batchsize=batch_size, resolution = img_size ) test_dataloader = dataloaderTriplet(datatestroot, data_transforms, num_triplets = num_triplets_train, batchsize = batch_size, resolution= img_size ) # Instantiate model model = Resnet18Triplet(embedding_dimension=embedding_dimension,pretrained=pretrained) if model_architecture == "resnet18": model = Resnet18Triplet( embedding_dimension=embedding_dimension, pretrained=pretrained ) elif model_architecture == "resnet34": model = Resnet34Triplet( embedding_dimension=embedding_dimension, pretrained=pretrained ) elif model_architecture == "resnet50": model = Resnet50Triplet( embedding_dimension=embedding_dimension, pretrained=pretrained ) elif model_architecture == "resnet101": model = Resnet101Triplet( embedding_dimension=embedding_dimension, pretrained=pretrained ) elif model_architecture == "inceptionresnetv2": model = InceptionResnetV2Triplet( embedding_dimension=embedding_dimension, pretrained=pretrained ) print("Using {} model architecture.".format(model_architecture)) # Load model to GPU or multiple GPUs if available flag_train_gpu = torch.cuda.is_available() flag_train_multi_gpu = False """ if flag_train_gpu and torch.cuda.device_count() > 1: model = nn.DataParallel(model) model.cuda() flag_train_multi_gpu = True print('Using multi-gpu training.') elif flag_train_gpu and torch.cuda.device_count() == 1: model.cuda() print('Using single-gpu training.') """ cuda0 = torch.device("cuda:0") model.to(cuda0) # Set optimizers if optimizer == "sgd": optimizer_model = torch.optim.SGD(model.parameters(), lr=learning_rate) elif optimizer == "adagrad": optimizer_model = torch.optim.Adagrad(model.parameters(), lr=learning_rate) elif optimizer == "rmsprop": optimizer_model = torch.optim.RMSprop(model.parameters(), lr=learning_rate) elif optimizer == "adam": optimizer_model = torch.optim.Adam(model.parameters(), lr=learning_rate) # Optionally resume from a checkpoint """ if resume_path: if os.path.isfile(resume_path): print("\nLoading checkpoint {} ...".format(resume_path)) checkpoint = torch.load(resume_path) start_epoch = checkpoint['epoch'] # In order to load state dict for optimizers correctly, model has to be loaded to gpu first if flag_train_multi_gpu: model.module.load_state_dict(checkpoint['model_state_dict']) else: model.load_state_dict(checkpoint['model_state_dict']) optimizer_model.load_state_dict(checkpoint['optimizer_model_state_dict']) print("\nCheckpoint loaded: start epoch from checkpoint = {}\nRunning for {} epochs.\n".format( start_epoch, epochs-start_epoch ) ) else: print("WARNING: No checkpoint found at {}!\nTraining from scratch.".format(resume_path)) """ # Start Training loop print("\nTraining using triplet loss on {} triplets starting for {} epochs:\n".format(num_triplets_train,epochs-start_epoch)) total_time_start = time.time() start_epoch = start_epoch end_epoch = start_epoch + epochs l2_distance = PairwiseDistance(2).to(cuda0) for epoch in range(start_epoch, end_epoch): epoch_time_start = time.time() triplet_loss_sum = 0 num_valid_for_training_triplets = 0 num_total_train_phase_triplets = 0 num_total_test_phase_triplets = 0 num_right_train_phase_triplets = 0 num_right_test_phase_triplets = 0 for phase in ['train', 'test']: if phase == 'train': model.train() dataloader = train_dataloader else : with torch.no_grad(): model.eval() dataloader = test_dataloader for batch_idx, (batch_sample) in enumerate(dataloader): if phase == 'test': with torch.no_grad(): anc_img = batch_sample[0].to(cuda0) pos_img = batch_sample[1].to(cuda0) neg_img = batch_sample[2].to(cuda0) # Forward pass - compute embedding anc_embedding = model(anc_img) pos_embedding = model(pos_img) neg_embedding = model(neg_img) # Forward pass - choose hard negatives only for training pos_dist = l2_distance.forward(anc_embedding, pos_embedding).cpu() neg_dist = l2_distance.forward(anc_embedding, neg_embedding).cpu() all = (neg_dist - pos_dist < margin).numpy().flatten() hard_triplets = np.where(all == 1) num_total_test_phase_triplets += len(all) num_right_test_phase_triplets += len(all)-len(hard_triplets[0]) if phase == 'train': anc_img = batch_sample[0].to(cuda0) pos_img = batch_sample[1].to(cuda0) neg_img = batch_sample[2].to(cuda0) # Forward pass - compute embedding anc_embedding = model(anc_img) pos_embedding = model(pos_img) neg_embedding = model(neg_img) # Forward pass - choose hard negatives only for training pos_dist = l2_distance.forward(anc_embedding, pos_embedding).cpu() neg_dist = l2_distance.forward(anc_embedding, neg_embedding).cpu() all = (neg_dist - pos_dist < margin).numpy().flatten() hard_triplets = np.where(all == 1) num_total_train_phase_triplets += len(all) num_right_train_phase_triplets += len(all)-len(hard_triplets[0]) if len(hard_triplets[0]) == 0: continue anc_hard_embedding = anc_embedding[hard_triplets].to(cuda0) pos_hard_embedding = pos_embedding[hard_triplets].to(cuda0) neg_hard_embedding = neg_embedding[hard_triplets].to(cuda0) # Calculate triplet loss triplet_loss = TripletLoss(margin=margin).forward(anchor=anc_hard_embedding,positive=pos_hard_embedding,negative=neg_hard_embedding).to(cuda0) # Calculating loss triplet_loss_sum += triplet_loss.item() if phase == 'train': num_valid_for_training_triplets += len(anc_hard_embedding) optimizer_model.zero_grad() if phase == 'train': # Backward pass triplet_loss.backward() optimizer_model.step() # Model only trains on hard negative triplets avg_triplet_loss = 0 if (num_valid_for_training_triplets == 0) else triplet_loss_sum / num_valid_for_training_triplets epoch_time_end = time.time() if phase == 'train': accuracy = (1.0 * num_right_train_phase_triplets )/(1.0 *num_total_train_phase_triplets) if phase == 'test': accuracy = (1.0 * num_right_test_phase_triplets) / (1.0 * num_total_test_phase_triplets) # Print training statistics and add to log print('Epoch {} Phase {}:\tAverage Triplet Loss: {:.4f}\tEpoch Time: {:.3f} hours\tNumber of valid training triplets : {}\tAccuracy : {}'.format(epoch+1,phase, avg_triplet_loss,(epoch_time_end - epoch_time_start)/3600,num_valid_for_training_triplets, accuracy)) with open('logs/{}_log_triplet_{}.txt'.format(model_architecture, phase), 'a') as f: val_list = [epoch+1,phase,avg_triplet_loss,num_valid_for_training_triplets, accuracy] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') # Save model checkpoint state = { 'epoch': epoch+1, 'embedding_dimension': embedding_dimension, 'batch_size_training': batch_size, 'model_state_dict': model.state_dict(), 'model_architecture': model_architecture, 'optimizer_model_state_dict': optimizer_model.state_dict() } # For storing data parallel model's state dictionary without 'module' parameter if flag_train_multi_gpu: state['model_state_dict'] = model.module.state_dict() # Save model checkpoint torch.save(state, 'Model_training_checkpoints/model_{}_triplet_epoch_{}.pt'.format(model_architecture, epoch+1)) # Training loop end total_time_end = time.time() total_time_elapsed = total_time_end - total_time_start print("\nTraining finished: total time elapsed: {:.2f} hours.".format(total_time_elapsed/3600))
def train_triplet(start_epoch, end_epoch, epochs, train_dataloader, lfw_dataloader, lfw_validation_epoch_interval, model, model_architecture, optimizer_model, embedding_dimension, batch_size, margin, flag_train_multi_gpu, optimizer, learning_rate): for epoch in range(start_epoch, end_epoch): flag_validate_lfw = (epoch + 1) % lfw_validation_epoch_interval == 0 or ( epoch + 1) % epochs == 0 triplet_loss_sum = 0 num_valid_training_triplets = 0 l2_distance = PairwiseDistance(2) # Training pass model.train() progress_bar = enumerate(tqdm(train_dataloader)) for batch_idx, (batch_sample) in progress_bar: # Forward pass - compute embeddings anc_imgs = batch_sample['anc_img'] pos_imgs = batch_sample['pos_img'] neg_imgs = batch_sample['neg_img'] anc_embeddings, pos_embeddings, neg_embeddings, model, optimizer_model, flag_use_cpu = forward_pass( anc_imgs=anc_imgs, pos_imgs=pos_imgs, neg_imgs=neg_imgs, model=model, optimizer_model=optimizer_model, batch_idx=batch_idx, optimizer=optimizer, learning_rate=learning_rate, use_cpu=False) # Forward pass - choose hard negatives only for training pos_dists = l2_distance.forward(anc_embeddings, pos_embeddings) neg_dists = l2_distance.forward(anc_embeddings, neg_embeddings) all = (neg_dists - pos_dists < margin).cpu().numpy().flatten() hard_triplets = np.where(all == 1) if len(hard_triplets[0]) == 0: continue anc_hard_embeddings = anc_embeddings[hard_triplets] pos_hard_embeddings = pos_embeddings[hard_triplets] neg_hard_embeddings = neg_embeddings[hard_triplets] # Calculate triplet loss triplet_loss = TripletLoss(margin=margin).forward( anchor=anc_hard_embeddings, positive=pos_hard_embeddings, negative=neg_hard_embeddings) # Calculating loss triplet_loss_sum += triplet_loss.item() num_valid_training_triplets += len(anc_hard_embeddings) # Backward pass optimizer_model.zero_grad() triplet_loss.backward() optimizer_model.step() # Load model and optimizer back to GPU if CUDA Out of Memory Exception occured and model and optimizer # were switched to CPU if flag_use_cpu: # According to https://github.com/pytorch/pytorch/issues/2830#issuecomment-336183179 # In order for the optimizer to keep training the model after changing to a different type or device, # optimizers have to be recreated, 'load_state_dict' can be used to restore the state from a # previous copy. As such, the optimizer state dict will be saved first and then reloaded when # the model's device is changed. # Load back to CUDA torch.save( optimizer_model.state_dict(), 'model_training_checkpoints/out_of_memory_optimizer_checkpoint/optimizer_checkpoint.pt' ) model = model.cuda() optimizer_model = set_optimizer(optimizer=optimizer, model=model, learning_rate=learning_rate) optimizer_model.load_state_dict( torch.load( 'model_training_checkpoints/out_of_memory_optimizer_checkpoint/optimizer_checkpoint.pt' )) # Copied from https://github.com/pytorch/pytorch/issues/2830#issuecomment-336194949 # No optimizer.cuda() available, this is the way to make an optimizer loaded with cpu tensors load # with cuda tensors. for state in optimizer_model.state.values(): for k, v in state.items(): if torch.is_tensor(v): state[k] = v.cuda() # Ensure model is correctly set to be trainable model.train() # Model only trains on hard negative triplets avg_triplet_loss = 0 if ( num_valid_training_triplets == 0) else triplet_loss_sum / num_valid_training_triplets # Print training statistics and add to log print( 'Epoch {}:\tAverage Triplet Loss: {:.4f}\tNumber of valid training triplets in epoch: {}' .format(epoch + 1, avg_triplet_loss, num_valid_training_triplets)) with open('logs/{}_log_triplet.txt'.format(model_architecture), 'a') as f: val_list = [ epoch + 1, avg_triplet_loss, num_valid_training_triplets ] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') try: # Plot Triplet losses plot plot_triplet_losses( log_dir="logs/{}_log_triplet.txt".format(model_architecture), epochs=epochs, figure_name="plots/triplet_losses_{}.png".format( model_architecture)) except Exception as e: print(e) # Evaluation pass on LFW dataset if flag_validate_lfw: best_distances = validate_lfw( model=model, lfw_dataloader=lfw_dataloader, model_architecture=model_architecture, epoch=epoch, epochs=epochs) # Save model checkpoint state = { 'epoch': epoch + 1, 'embedding_dimension': embedding_dimension, 'batch_size_training': batch_size, 'model_state_dict': model.state_dict(), 'model_architecture': model_architecture, 'optimizer_model_state_dict': optimizer_model.state_dict() } # For storing data parallel model's state dictionary without 'module' parameter if flag_train_multi_gpu: state['model_state_dict'] = model.module.state_dict() # For storing best euclidean distance threshold during LFW validation if flag_validate_lfw: state['best_distance_threshold'] = np.mean(best_distances) # Save model checkpoint torch.save( state, 'model_training_checkpoints/model_{}_triplet_epoch_{}.pt'.format( model_architecture, epoch + 1))
def __init__(self): self.pdist = PairwiseDistance(2) self.loss_fn = nn.MSELoss() self.b = config["threshold loss"]
def __init__(self, network, learning_rate): self.network = network self.learning_rate = learning_rate self.optimizer = torch.optim.Adam(self.network.parameters(), lr=self.learning_rate) self.criterion = nn.TripletMarginLoss(margin=0.2) self.l2_dist = PairwiseDistance(2)
def __init__(self, margin=1.): super(ITripletLoss, self).__init__() self.margin = margin self.pair_loss = PairwiseDistance(2)
dataset_loader = DataLoader(dataset, batch_size=10, shuffle=False, num_workers=4) ImageIds = [] for x in dataset.imgs: ImageIds.append(x[0][-20:-4]) ImageId = pd.DataFrame(ImageIds) ImageId.columns = ['ImageId'] labels = ImageId.merge(labels_init, on='ImageId') # In[ ]: EPSILON = 8 / 255 * (1 - -1) l2dist = PairwiseDistance(2) criterion_none = nn.CrossEntropyLoss(reduction='none') criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4) # attacks bim_attack = Attack(net, F.cross_entropy) cw_attack = cw.L2Adversary(targeted=False, confidence=0.9, search_steps=10, box=(-1, 1), optimizer_lr=0.001) # In[ ]:
def main(): dataroot = args.dataroot lfw_dataroot = args.lfw training_dataset_csv_path = args.training_dataset_csv_path epochs = args.epochs iterations_per_epoch = args.iterations_per_epoch model_architecture = args.model_architecture pretrained = args.pretrained embedding_dimension = args.embedding_dimension num_human_identities_per_batch = args.num_human_identities_per_batch batch_size = args.batch_size lfw_batch_size = args.lfw_batch_size resume_path = args.resume_path num_workers = args.num_workers optimizer = args.optimizer learning_rate = args.learning_rate margin = args.margin image_size = args.image_size use_semihard_negatives = args.use_semihard_negatives training_triplets_path = args.training_triplets_path flag_training_triplets_path = False start_epoch = 0 if training_triplets_path is not None: flag_training_triplets_path = True # Load triplets file for the first training epoch # Define image data pre-processing transforms # ToTensor() normalizes pixel values between [0, 1] # Normalize(mean=[0.6071, 0.4609, 0.3944], std=[0.2457, 0.2175, 0.2129]) according to the calculated glint360k # dataset with tightly-cropped faces dataset RGB channels' mean and std values by # calculate_glint360k_rgb_mean_std.py in 'datasets' folder. data_transforms = transforms.Compose([ transforms.Resize(size=image_size), transforms.RandomHorizontalFlip(), transforms.RandomRotation(degrees=5), transforms.ToTensor(), transforms.Normalize(mean=[0.6071, 0.4609, 0.3944], std=[0.2457, 0.2175, 0.2129]) ]) lfw_transforms = transforms.Compose([ transforms.Resize(size=image_size), transforms.ToTensor(), transforms.Normalize(mean=[0.6071, 0.4609, 0.3944], std=[0.2457, 0.2175, 0.2129]) ]) lfw_dataloader = torch.utils.data.DataLoader(dataset=LFWDataset( dir=lfw_dataroot, pairs_path='datasets/LFW_pairs.txt', transform=lfw_transforms), batch_size=lfw_batch_size, num_workers=num_workers, shuffle=False) # Instantiate model model = set_model_architecture(model_architecture=model_architecture, pretrained=pretrained, embedding_dimension=embedding_dimension) # Load model to GPU or multiple GPUs if available model, flag_train_multi_gpu = set_model_gpu_mode(model) # Set optimizer optimizer_model = set_optimizer(optimizer=optimizer, model=model, learning_rate=learning_rate) # Resume from a model checkpoint if resume_path: if os.path.isfile(resume_path): print("Loading checkpoint {} ...".format(resume_path)) checkpoint = torch.load(resume_path) start_epoch = checkpoint['epoch'] + 1 optimizer_model.load_state_dict( checkpoint['optimizer_model_state_dict']) # In order to load state dict for optimizers correctly, model has to be loaded to gpu first if flag_train_multi_gpu: model.module.load_state_dict(checkpoint['model_state_dict']) else: model.load_state_dict(checkpoint['model_state_dict']) print("Checkpoint loaded: start epoch from checkpoint = {}".format( start_epoch)) else: print( "WARNING: No checkpoint found at {}!\nTraining from scratch.". format(resume_path)) if use_semihard_negatives: print("Using Semi-Hard negative triplet selection!") else: print("Using Hard negative triplet selection!") start_epoch = start_epoch print("Training using triplet loss starting for {} epochs:\n".format( epochs - start_epoch)) for epoch in range(start_epoch, epochs): num_valid_training_triplets = 0 l2_distance = PairwiseDistance(p=2) _training_triplets_path = None if flag_training_triplets_path: _training_triplets_path = training_triplets_path flag_training_triplets_path = False # Only load triplets file for the first epoch # Re-instantiate training dataloader to generate a triplet list for this training epoch train_dataloader = torch.utils.data.DataLoader( dataset=TripletFaceDataset( root_dir=dataroot, training_dataset_csv_path=training_dataset_csv_path, num_triplets=iterations_per_epoch * batch_size, num_human_identities_per_batch=num_human_identities_per_batch, triplet_batch_size=batch_size, epoch=epoch, training_triplets_path=_training_triplets_path, transform=data_transforms), batch_size=batch_size, num_workers=num_workers, shuffle= False # Shuffling for triplets with set amount of human identities per batch is not required ) # Training pass model.train() progress_bar = enumerate(tqdm(train_dataloader)) for batch_idx, (batch_sample) in progress_bar: # Forward pass - compute embeddings anc_imgs = batch_sample['anc_img'] pos_imgs = batch_sample['pos_img'] neg_imgs = batch_sample['neg_img'] # Concatenate the input images into one tensor because doing multiple forward passes would create # weird GPU memory allocation behaviours later on during training which would cause GPU Out of Memory # issues all_imgs = torch.cat( (anc_imgs, pos_imgs, neg_imgs)) # Must be a tuple of Torch Tensors anc_embeddings, pos_embeddings, neg_embeddings, model = forward_pass( imgs=all_imgs, model=model, batch_size=batch_size) pos_dists = l2_distance.forward(anc_embeddings, pos_embeddings) neg_dists = l2_distance.forward(anc_embeddings, neg_embeddings) if use_semihard_negatives: # Semi-Hard Negative triplet selection # (negative_distance - positive_distance < margin) AND (positive_distance < negative_distance) # Based on: https://github.com/davidsandberg/facenet/blob/master/src/train_tripletloss.py#L295 first_condition = (neg_dists - pos_dists < margin).cpu().numpy().flatten() second_condition = (pos_dists < neg_dists).cpu().numpy().flatten() all = (np.logical_and(first_condition, second_condition)) valid_triplets = np.where(all == 1) else: # Hard Negative triplet selection # (negative_distance - positive_distance < margin) # Based on: https://github.com/davidsandberg/facenet/blob/master/src/train_tripletloss.py#L296 all = (neg_dists - pos_dists < margin).cpu().numpy().flatten() valid_triplets = np.where(all == 1) anc_valid_embeddings = anc_embeddings[valid_triplets] pos_valid_embeddings = pos_embeddings[valid_triplets] neg_valid_embeddings = neg_embeddings[valid_triplets] # Calculate triplet loss triplet_loss = TripletLoss(margin=margin).forward( anchor=anc_valid_embeddings, positive=pos_valid_embeddings, negative=neg_valid_embeddings) # Calculating number of triplets that met the triplet selection method during the epoch num_valid_training_triplets += len(anc_valid_embeddings) # Backward pass optimizer_model.zero_grad() triplet_loss.backward() optimizer_model.step() # Print training statistics for epoch and add to log print( 'Epoch {}:\tNumber of valid training triplets in epoch: {}'.format( epoch, num_valid_training_triplets)) with open('logs/{}_log_triplet.txt'.format(model_architecture), 'a') as f: val_list = [epoch, num_valid_training_triplets] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') # Evaluation pass on LFW dataset best_distances = validate_lfw(model=model, lfw_dataloader=lfw_dataloader, model_architecture=model_architecture, epoch=epoch) # Save model checkpoint state = { 'epoch': epoch, 'embedding_dimension': embedding_dimension, 'batch_size_training': batch_size, 'model_state_dict': model.state_dict(), 'model_architecture': model_architecture, 'optimizer_model_state_dict': optimizer_model.state_dict(), 'best_distance_threshold': np.mean(best_distances) } # For storing data parallel model's state dictionary without 'module' parameter if flag_train_multi_gpu: state['model_state_dict'] = model.module.state_dict() # Save model checkpoint torch.save( state, 'model_training_checkpoints/model_{}_triplet_epoch_{}.pt'.format( model_architecture, epoch))
def train_triplet(start_epoch, end_epoch, epochs, train_dataloader, lfw_dataloader, lfw_validation_epoch_interval, model, model_architecture, optimizer_model, embedding_dimension, batch_size, margin, flag_train_multi_gpu, optimizer, learning_rate, use_semihard_negatives): for epoch in range(start_epoch, end_epoch): flag_validate_lfw = (epoch + 1) % lfw_validation_epoch_interval == 0 or ( epoch + 1) % epochs == 0 triplet_loss_sum = 0 num_valid_training_triplets = 0 l2_distance = PairwiseDistance(p=2) # Training pass model.train() progress_bar = enumerate(tqdm(train_dataloader)) for batch_idx, (batch_sample) in progress_bar: # Skip last iteration to avoid the problem of having different number of tensors while calculating # pairwise distances (sizes of tensors must be the same for pairwise distance calculation) if batch_idx + 1 == len(train_dataloader): continue # Forward pass - compute embeddings anc_imgs = batch_sample['anc_img'] pos_imgs = batch_sample['pos_img'] neg_imgs = batch_sample['neg_img'] # Concatenate the input images into one tensor because doing multiple forward passes would create # weird GPU memory allocation behaviours later on during training which would cause GPU Out of Memory # issues all_imgs = torch.cat( (anc_imgs, pos_imgs, neg_imgs)) # Must be a tuple of Torch Tensors anc_embeddings, pos_embeddings, neg_embeddings, model, optimizer_model, flag_use_cpu = forward_pass( imgs=all_imgs, model=model, optimizer_model=optimizer_model, batch_idx=batch_idx, optimizer=optimizer, learning_rate=learning_rate, batch_size=batch_size, use_cpu=False) pos_dists = l2_distance.forward(anc_embeddings, pos_embeddings) neg_dists = l2_distance.forward(anc_embeddings, neg_embeddings) if use_semihard_negatives: # Semi-Hard Negative triplet selection # (negative_distance - positive_distance < margin) AND (positive_distance < negative_distance) # Based on: https://github.com/davidsandberg/facenet/blob/master/src/train_tripletloss.py#L295 first_condition = (neg_dists - pos_dists < margin).cpu().numpy().flatten() second_condition = (pos_dists < neg_dists).cpu().numpy().flatten() all = (np.logical_and(first_condition, second_condition)) semihard_negative_triplets = np.where(all == 1) if len(semihard_negative_triplets[0]) == 0: continue anc_valid_embeddings = anc_embeddings[ semihard_negative_triplets] pos_valid_embeddings = pos_embeddings[ semihard_negative_triplets] neg_valid_embeddings = neg_embeddings[ semihard_negative_triplets] del anc_embeddings, pos_embeddings, neg_embeddings, pos_dists, neg_dists gc.collect() else: # Hard Negative triplet selection # (negative_distance - positive_distance < margin) # Based on: https://github.com/davidsandberg/facenet/blob/master/src/train_tripletloss.py#L296 all = (neg_dists - pos_dists < margin).cpu().numpy().flatten() hard_negative_triplets = np.where(all == 1) if len(hard_negative_triplets[0]) == 0: continue anc_valid_embeddings = anc_embeddings[hard_negative_triplets] pos_valid_embeddings = pos_embeddings[hard_negative_triplets] neg_valid_embeddings = neg_embeddings[hard_negative_triplets] del anc_embeddings, pos_embeddings, neg_embeddings, pos_dists, neg_dists gc.collect() # Calculate triplet loss triplet_loss = TripletLoss(margin=margin).forward( anchor=anc_valid_embeddings, positive=pos_valid_embeddings, negative=neg_valid_embeddings) # Calculating loss and number of triplets that met the triplet selection method during the epoch triplet_loss_sum += triplet_loss.item() num_valid_training_triplets += len(anc_valid_embeddings) # Backward pass optimizer_model.zero_grad() triplet_loss.backward() optimizer_model.step() # Load model and optimizer back to GPU if CUDA Out of Memory Exception occurred and model and optimizer # were switched to CPU if flag_use_cpu: # According to https://github.com/pytorch/pytorch/issues/2830#issuecomment-336183179 # In order for the optimizer to keep training the model after changing to a different type or device, # optimizers have to be recreated, 'load_state_dict' can be used to restore the state from a # previous copy. As such, the optimizer state dict will be saved first and then reloaded when # the model's device is changed. torch.cuda.empty_cache() # Print number of valid triplets (troubleshooting out of memory causes) print("Number of valid triplets during OOM iteration = {}". format(len(anc_valid_embeddings))) torch.save( optimizer_model.state_dict(), 'model_training_checkpoints/out_of_memory_optimizer_checkpoint/optimizer_checkpoint.pt' ) # Load back to CUDA model.cuda() optimizer_model = set_optimizer(optimizer=optimizer, model=model, learning_rate=learning_rate) optimizer_model.load_state_dict( torch.load( 'model_training_checkpoints/out_of_memory_optimizer_checkpoint/optimizer_checkpoint.pt' )) # Copied from https://github.com/pytorch/pytorch/issues/2830#issuecomment-336194949 # No optimizer.cuda() available, this is the way to make an optimizer loaded with cpu tensors load # with cuda tensors. for state in optimizer_model.state.values(): for k, v in state.items(): if torch.is_tensor(v): state[k] = v.cuda() # Clear some memory at end of training iteration del triplet_loss, anc_valid_embeddings, pos_valid_embeddings, neg_valid_embeddings gc.collect() # Model only trains on triplets that fit the triplet selection method avg_triplet_loss = 0 if ( num_valid_training_triplets == 0) else triplet_loss_sum / num_valid_training_triplets # Print training statistics and add to log print( 'Epoch {}:\tAverage Triplet Loss: {:.4f}\tNumber of valid training triplets in epoch: {}' .format(epoch + 1, avg_triplet_loss, num_valid_training_triplets)) with open('logs/{}_log_triplet.txt'.format(model_architecture), 'a') as f: val_list = [ epoch + 1, avg_triplet_loss, num_valid_training_triplets ] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') try: # Plot Triplet losses plot plot_triplet_losses( log_dir="logs/{}_log_triplet.txt".format(model_architecture), epochs=epochs, figure_name="plots/triplet_losses_{}.png".format( model_architecture)) except Exception as e: print(e) # Evaluation pass on LFW dataset if flag_validate_lfw: best_distances = validate_lfw( model=model, lfw_dataloader=lfw_dataloader, model_architecture=model_architecture, epoch=epoch, epochs=epochs) # Save model checkpoint state = { 'epoch': epoch + 1, 'embedding_dimension': embedding_dimension, 'batch_size_training': batch_size, 'model_state_dict': model.state_dict(), 'model_architecture': model_architecture, 'optimizer_model_state_dict': optimizer_model.state_dict() } # For storing data parallel model's state dictionary without 'module' parameter if flag_train_multi_gpu: state['model_state_dict'] = model.module.state_dict() # For storing best euclidean distance threshold during LFW validation if flag_validate_lfw: state['best_distance_threshold'] = np.mean(best_distances) # Save model checkpoint torch.save( state, 'model_training_checkpoints/model_{}_triplet_epoch_{}.pt'.format( model_architecture, epoch + 1))
def validate_lfw(model, lfw_dataloader, model_architecture, epoch, epochs): model.eval() with torch.no_grad(): l2_distance = PairwiseDistance(2).cuda() distances, labels = [], [] print("Validating on LFW! ...") progress_bar = enumerate(tqdm(lfw_dataloader)) for batch_index, (data_a, data_b, label) in progress_bar: data_a, data_b, label = data_a.cuda(), data_b.cuda(), label.cuda() output_a, output_b = model(data_a), model(data_b) distance = l2_distance.forward(output_a, output_b) # Euclidean distance distances.append(distance.cpu().detach().numpy()) labels.append(label.cpu().detach().numpy()) labels = np.array([sublabel for label in labels for sublabel in label]) distances = np.array( [subdist for distance in distances for subdist in distance]) true_positive_rate, false_positive_rate, precision, recall, accuracy, roc_auc, best_distances, \ tar, far = evaluate_lfw( distances=distances, labels=labels, far_target=1e-3 ) # Print statistics and add to log print( "Accuracy on LFW: {:.4f}+-{:.4f}\tPrecision {:.4f}+-{:.4f}\tRecall {:.4f}+-{:.4f}\t" "ROC Area Under Curve: {:.4f}\tBest distance threshold: {:.2f}+-{:.2f}\t" "TAR: {:.4f}+-{:.4f} @ FAR: {:.4f}".format(np.mean(accuracy), np.std(accuracy), np.mean(precision), np.std(precision), np.mean(recall), np.std(recall), roc_auc, np.mean(best_distances), np.std(best_distances), np.mean(tar), np.std(tar), np.mean(far))) with open('logs/lfw_{}_log_triplet.txt'.format(model_architecture), 'a') as f: val_list = [ epoch + 1, np.mean(accuracy), np.std(accuracy), np.mean(precision), np.std(precision), np.mean(recall), np.std(recall), roc_auc, np.mean(best_distances), np.std(best_distances), np.mean(tar) ] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') try: # Plot ROC curve plot_roc_lfw( false_positive_rate=false_positive_rate, true_positive_rate=true_positive_rate, figure_name="plots/roc_plots/roc_{}_epoch_{}_triplet.png".format( model_architecture, epoch + 1)) # Plot LFW accuracies plot plot_accuracy_lfw( log_dir="logs/lfw_{}_log_triplet.txt".format(model_architecture), epochs=epochs, figure_name="plots/lfw_accuracies_{}_triplet.png".format( model_architecture)) except Exception as e: print(e) return best_distances
def train_triplet(start_epoch, end_epoch, epochs, train_dataloader, lfw_dataloader, lfw_validation_epoch_interval, model, model_architecture, optimizer_model, embedding_dimension, batch_size, margin, flag_train_multi_gpu): for epoch in range(start_epoch, end_epoch): flag_validate_lfw = (epoch + 1) % lfw_validation_epoch_interval == 0 or ( epoch + 1) % epochs == 0 triplet_loss_sum = 0 num_valid_training_triplets = 0 l2_distance = PairwiseDistance(2).cuda() # Training pass model.train() progress_bar = enumerate(tqdm(train_dataloader)) for batch_idx, (batch_sample) in progress_bar: anc_img = batch_sample['anc_img'].cuda() pos_img = batch_sample['pos_img'].cuda() neg_img = batch_sample['neg_img'].cuda() # Forward pass - compute embeddings anc_embedding, pos_embedding, neg_embedding = model( anc_img), model(pos_img), model(neg_img) # Forward pass - choose hard negatives only for training pos_dist = l2_distance.forward(anc_embedding, pos_embedding) neg_dist = l2_distance.forward(anc_embedding, neg_embedding) all = (neg_dist - pos_dist < margin).cpu().numpy().flatten() hard_triplets = np.where(all == 1) if len(hard_triplets[0]) == 0: continue anc_hard_embedding = anc_embedding[hard_triplets].cuda() pos_hard_embedding = pos_embedding[hard_triplets].cuda() neg_hard_embedding = neg_embedding[hard_triplets].cuda() # Calculate triplet loss triplet_loss = TripletLoss(margin=margin).forward( anchor=anc_hard_embedding, positive=pos_hard_embedding, negative=neg_hard_embedding).cuda() # Calculating loss triplet_loss_sum += triplet_loss.item() num_valid_training_triplets += len(anc_hard_embedding) # Backward pass optimizer_model.zero_grad() triplet_loss.backward() optimizer_model.step() # Model only trains on hard negative triplets avg_triplet_loss = 0 if ( num_valid_training_triplets == 0) else triplet_loss_sum / num_valid_training_triplets # Print training statistics and add to log print( 'Epoch {}:\tAverage Triplet Loss: {:.4f}\tNumber of valid training triplets in epoch: {}' .format(epoch + 1, avg_triplet_loss, num_valid_training_triplets)) with open('logs/{}_log_triplet.txt'.format(model_architecture), 'a') as f: val_list = [ epoch + 1, avg_triplet_loss, num_valid_training_triplets ] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') try: # Plot Triplet losses plot plot_triplet_losses( log_dir="logs/{}_log_triplet.txt".format(model_architecture), epochs=epochs, figure_name="plots/triplet_losses_{}.png".format( model_architecture)) except Exception as e: print(e) # Evaluation pass on LFW dataset if flag_validate_lfw: best_distances = validate_lfw( model=model, lfw_dataloader=lfw_dataloader, model_architecture=model_architecture, epoch=epoch, epochs=epochs) # Save model checkpoint state = { 'epoch': epoch + 1, 'embedding_dimension': embedding_dimension, 'batch_size_training': batch_size, 'model_state_dict': model.state_dict(), 'model_architecture': model_architecture, 'optimizer_model_state_dict': optimizer_model.state_dict() } # For storing data parallel model's state dictionary without 'module' parameter if flag_train_multi_gpu: state['model_state_dict'] = model.module.state_dict() # For storing best euclidean distance threshold during LFW validation if flag_validate_lfw: state['best_distance_threshold'] = np.mean(best_distances) # Save model checkpoint torch.save( state, 'Model_training_checkpoints/model_{}_triplet_epoch_{}.pt'.format( model_architecture, epoch + 1))
# 模型加载 model = Resnet18Triplet(pretrained=False, embedding_dimension=128) if torch.cuda.is_available(): model.cuda() print('Using single-gpu testing.') model_pathi = "../week21/Model_training_checkpoints/model_resnet18_triplet_epoch_603.pt" if os.path.exists(model_pathi): model_state = torch.load(model_pathi) model.load_state_dict(model_state['model_state_dict']) start_epoch = model_state['epoch'] print('loaded %s' % model_pathi) else: print('不存在预训练模型!') l2_distance = PairwiseDistance(2) with torch.no_grad(): # 不传梯度了 distances, labels = [], [] progress_bar = enumerate(tqdm(test_dataloader)) # track bar for batch_index, (data_a, data_b, label) in progress_bar: #for batch_index, (data_a, data_b, label) in enumerate(test_dataloader): # data_a, data_b, label这仨是一批的矩阵 data_a = data_a.cuda() data_b = data_b.cuda() label = label.cuda() output_a, output_b = model(data_a), model(data_b) output_a = torch.div(output_a, torch.norm(output_a)) output_b = torch.div(output_b, torch.norm(output_b)) distance = l2_distance.forward(output_a, output_b) # 列表里套矩阵 labels.append(label.cpu().detach().numpy())
def main(): dataroot = args.dataroot lfw_dataroot = args.lfw lfw_batch_size = args.lfw_batch_size lfw_validation_epoch_interval = args.lfw_validation_epoch_interval model_architecture = args.model epochs = args.epochs resume_path = args.resume_path batch_size = args.batch_size num_workers = args.num_workers validation_dataset_split_ratio = args.valid_split embedding_dimension = args.embedding_dim pretrained = args.pretrained optimizer = args.optimizer learning_rate = args.lr learning_rate_center_loss = args.center_loss_lr center_loss_weight = args.center_loss_weight start_epoch = 0 # Define image data pre-processing transforms # ToTensor() normalizes pixel values between [0, 1] # Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) normalizes pixel values between [-1, 1] # Size 182x182 RGB image -> Center crop size 160x160 RGB image for more model generalization data_transforms = transforms.Compose([ transforms.RandomCrop(size=160), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # Size 160x160 RGB image lfw_transforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # Load the dataset dataset = torchvision.datasets.ImageFolder(root=dataroot, transform=data_transforms) # Subset the dataset into training and validation datasets num_classes = len(dataset.classes) print("\nNumber of classes in dataset: {}".format(num_classes)) num_validation = int(num_classes * validation_dataset_split_ratio) num_train = num_classes - num_validation indices = list(range(num_classes)) np.random.seed(420) np.random.shuffle(indices) train_indices = indices[:num_train] validation_indices = indices[num_train:] train_dataset = Subset(dataset=dataset, indices=train_indices) validation_dataset = Subset(dataset=dataset, indices=validation_indices) print("Number of classes in training dataset: {}".format( len(train_dataset))) print("Number of classes in validation dataset: {}".format( len(validation_dataset))) # Define the dataloaders train_dataloader = DataLoader(dataset=train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False) validation_dataloader = DataLoader(dataset=validation_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False) lfw_dataloader = torch.utils.data.DataLoader(dataset=LFWDataset( dir=lfw_dataroot, pairs_path='datasets/LFW_pairs.txt', transform=lfw_transforms), batch_size=lfw_batch_size, num_workers=num_workers, shuffle=False) # Instantiate model if model_architecture == "resnet18": model = Resnet18Center(num_classes=num_classes, embedding_dimension=embedding_dimension, pretrained=pretrained) elif model_architecture == "resnet34": model = Resnet34Center(num_classes=num_classes, embedding_dimension=embedding_dimension, pretrained=pretrained) elif model_architecture == "resnet50": model = Resnet50Center(num_classes=num_classes, embedding_dimension=embedding_dimension, pretrained=pretrained) elif model_architecture == "resnet101": model = Resnet101Center(num_classes=num_classes, embedding_dimension=embedding_dimension, pretrained=pretrained) elif model_architecture == "inceptionresnetv2": model = InceptionResnetV2Center( num_classes=num_classes, embedding_dimension=embedding_dimension, pretrained=pretrained) print("\nUsing {} model architecture.".format(model_architecture)) # Load model to GPU or multiple GPUs if available flag_train_gpu = torch.cuda.is_available() flag_train_multi_gpu = False if flag_train_gpu and torch.cuda.device_count() > 1: model = nn.DataParallel(model) model.cuda() flag_train_multi_gpu = True print('Using multi-gpu training.') elif flag_train_gpu and torch.cuda.device_count() == 1: model.cuda() print('Using single-gpu training.') # Set loss functions criterion_crossentropy = nn.CrossEntropyLoss().cuda() criterion_centerloss = CenterLoss(num_classes=num_classes, feat_dim=embedding_dimension).cuda() # Set optimizers if optimizer == "sgd": optimizer_model = torch.optim.SGD(model.parameters(), lr=learning_rate) optimizer_centerloss = torch.optim.SGD( criterion_centerloss.parameters(), lr=learning_rate_center_loss) elif optimizer == "adagrad": optimizer_model = torch.optim.Adagrad(model.parameters(), lr=learning_rate) optimizer_centerloss = torch.optim.Adagrad( criterion_centerloss.parameters(), lr=learning_rate_center_loss) elif optimizer == "rmsprop": optimizer_model = torch.optim.RMSprop(model.parameters(), lr=learning_rate) optimizer_centerloss = torch.optim.RMSprop( criterion_centerloss.parameters(), lr=learning_rate_center_loss) elif optimizer == "adam": optimizer_model = torch.optim.Adam(model.parameters(), lr=learning_rate) optimizer_centerloss = torch.optim.Adam( criterion_centerloss.parameters(), lr=learning_rate_center_loss) # Set learning rate decay scheduler learning_rate_scheduler = optim.lr_scheduler.MultiStepLR( optimizer=optimizer_model, milestones=[150, 225], gamma=0.1) # Optionally resume from a checkpoint if resume_path: if os.path.isfile(resume_path): print("\nLoading checkpoint {} ...".format(resume_path)) checkpoint = torch.load(resume_path) start_epoch = checkpoint['epoch'] # In order to load state dict for optimizers correctly, model has to be loaded to gpu first if flag_train_multi_gpu: model.module.load_state_dict(checkpoint['model_state_dict']) else: model.load_state_dict(checkpoint['model_state_dict']) optimizer_model.load_state_dict( checkpoint['optimizer_model_state_dict']) optimizer_centerloss.load_state_dict( checkpoint['optimizer_centerloss_state_dict']) learning_rate_scheduler.load_state_dict( checkpoint['learning_rate_scheduler_state_dict']) print( "\nCheckpoint loaded: start epoch from checkpoint = {}\nRunning for {} epochs.\n" .format(start_epoch, epochs - start_epoch)) else: print( "WARNING: No checkpoint found at {}!\nTraining from scratch.". format(resume_path)) # Start Training loop print( "\nTraining using cross entropy loss with center loss starting for {} epochs:\n" .format(epochs - start_epoch)) total_time_start = time.time() start_epoch = start_epoch end_epoch = start_epoch + epochs for epoch in range(start_epoch, end_epoch): epoch_time_start = time.time() flag_validate_lfw = (epoch + 1) % lfw_validation_epoch_interval == 0 or ( epoch + 1) % epochs == 0 train_loss_sum = 0 validation_loss_sum = 0 # Training the model model.train() learning_rate_scheduler.step() progress_bar = enumerate(tqdm(train_dataloader)) for batch_index, (data, labels) in progress_bar: data, labels = data.cuda(), labels.cuda() # Forward pass if flag_train_multi_gpu: embedding, logits = model.module.forward_training(data) else: embedding, logits = model.forward_training(data) # Calculate losses cross_entropy_loss = criterion_crossentropy( logits.cuda(), labels.cuda()) center_loss = criterion_centerloss(embedding, labels) loss = (center_loss * center_loss_weight) + cross_entropy_loss # Backward pass optimizer_centerloss.zero_grad() optimizer_model.zero_grad() loss.backward() optimizer_centerloss.step() optimizer_model.step() # Remove center_loss_weight impact on the learning of center vectors for param in criterion_centerloss.parameters(): param.grad.data *= (1. / center_loss_weight) # Update training loss sum train_loss_sum += loss.item() * data.size(0) # Validating the model model.eval() correct, total = 0, 0 with torch.no_grad(): progress_bar = enumerate(tqdm(validation_dataloader)) for batch_index, (data, labels) in progress_bar: data, labels = data.cuda(), labels.cuda() # Forward pass if flag_train_multi_gpu: embedding, logits = model.module.forward_training(data) else: embedding, logits = model.forward_training(data) # Calculate losses cross_entropy_loss = criterion_crossentropy( logits.cuda(), labels.cuda()) center_loss = criterion_centerloss(embedding, labels) loss = (center_loss * center_loss_weight) + cross_entropy_loss # Update average validation loss validation_loss_sum += loss.item() * data.size(0) # Calculate training performance metrics predictions = logits.data.max(1)[1] total += labels.size(0) correct += (predictions == labels.data).sum() # Calculate average losses in epoch avg_train_loss = train_loss_sum / len(train_dataloader.dataset) avg_validation_loss = validation_loss_sum / len( validation_dataloader.dataset) # Calculate training performance statistics in epoch classification_accuracy = correct * 100. / total classification_error = 100. - classification_accuracy epoch_time_end = time.time() # Print training and validation statistics and add to log print( 'Epoch {}:\t Average Training Loss: {:.4f}\tAverage Validation Loss: {:.4f}\tClassification Accuracy: {:.2f}%\tClassification Error: {:.2f}%\tEpoch Time: {:.3f} hours' .format(epoch + 1, avg_train_loss, avg_validation_loss, classification_accuracy, classification_error, (epoch_time_end - epoch_time_start) / 3600)) with open('logs/{}_log_center.txt'.format(model_architecture), 'a') as f: val_list = [ epoch + 1, avg_train_loss, avg_validation_loss, classification_accuracy.item(), classification_error.item() ] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') try: # Plot plot for Cross Entropy Loss and Center Loss on training and validation sets plot_training_validation_losses_center( log_dir="logs/{}_log_center.txt".format(model_architecture), epochs=epochs, figure_name="plots/training_validation_losses_{}_center.png". format(model_architecture)) except Exception as e: print(e) # Validating on LFW dataset using KFold based on Euclidean distance metric if flag_validate_lfw: model.eval() with torch.no_grad(): l2_distance = PairwiseDistance(2).cuda() distances, labels = [], [] print("Validating on LFW! ...") progress_bar = enumerate(tqdm(lfw_dataloader)) for batch_index, (data_a, data_b, label) in progress_bar: data_a, data_b, label = data_a.cuda(), data_b.cuda( ), label.cuda() output_a, output_b = model(data_a), model(data_b) distance = l2_distance.forward( output_a, output_b) # Euclidean distance distances.append(distance.cpu().detach().numpy()) labels.append(label.cpu().detach().numpy()) labels = np.array( [sublabel for label in labels for sublabel in label]) distances = np.array([ subdist for distance in distances for subdist in distance ]) true_positive_rate, false_positive_rate, precision, recall, accuracy, roc_auc, best_distances, \ tar, far = evaluate_lfw( distances=distances, labels=labels ) # Print statistics and add to log print( "Accuracy on LFW: {:.4f}+-{:.4f}\tPrecision {:.4f}+-{:.4f}\tRecall {:.4f}+-{:.4f}\tROC Area Under Curve: {:.4f}\tBest distance threshold: {:.2f}+-{:.2f}\tTAR: {:.4f}+-{:.4f} @ FAR: {:.4f}" .format(np.mean(accuracy), np.std(accuracy), np.mean(precision), np.std(precision), np.mean(recall), np.std(recall), roc_auc, np.mean(best_distances), np.std(best_distances), np.mean(tar), np.std(tar), np.mean(far))) with open( 'logs/lfw_{}_log_center.txt'.format( model_architecture), 'a') as f: val_list = [ epoch + 1, np.mean(accuracy), np.std(accuracy), np.mean(precision), np.std(precision), np.mean(recall), np.std(recall), roc_auc, np.mean(best_distances), np.std(best_distances), np.mean(tar) ] log = '\t'.join(str(value) for value in val_list) f.writelines(log + '\n') try: # Plot ROC curve plot_roc_lfw( false_positive_rate=false_positive_rate, true_positive_rate=true_positive_rate, figure_name="plots/roc_plots/roc_{}_epoch_{}_center.png". format(model_architecture, epoch + 1)) # Plot LFW accuracies plot plot_accuracy_lfw( log_dir="logs/lfw_{}_log_center.txt".format( model_architecture), epochs=epochs, figure_name="plots/lfw_accuracies_{}_center.png".format( model_architecture)) except Exception as e: print(e) # Save model checkpoint state = { 'epoch': epoch + 1, 'num_classes': num_classes, 'embedding_dimension': embedding_dimension, 'batch_size_training': batch_size, 'model_state_dict': model.state_dict(), 'model_architecture': model_architecture, 'optimizer_model_state_dict': optimizer_model.state_dict(), 'optimizer_centerloss_state_dict': optimizer_centerloss.state_dict(), 'learning_rate_scheduler_state_dict': learning_rate_scheduler.state_dict() } # For storing data parallel model's state dictionary without 'module' parameter if flag_train_multi_gpu: state['model_state_dict'] = model.module.state_dict() # For storing best euclidean distance threshold during LFW validation if flag_validate_lfw: state['best_distance_threshold'] = np.mean(best_distances) # Save model checkpoint torch.save( state, 'Model_training_checkpoints/model_{}_center_epoch_{}.pt'.format( model_architecture, epoch + 1)) # Training loop end total_time_end = time.time() total_time_elapsed = total_time_end - total_time_start print("\nTraining finished: total time elapsed: {:.2f} hours.".format( total_time_elapsed / 3600))