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 __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 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)))
class TripletLoss(Function):
def __init__(self,margin):
super (TripletLoss,self).__init__()
self.margin= margin
self.pdist= PairwiseDistance(2)
def forward(self,anchor,positive,negative):
pos_dist = self.pdist.forward(anchor, positive)
neg_dist = self.pdist.forward(anchor,negative)
hinge_dist = torch.clamp(self.margin+pos_dist-neg_dist,min=0.0)
loss = torch.mean(hinge_dist)
return loss
class TripletLoss(torch.nn.Module):
def __init__(self, margin):
super(TripletLoss, self).__init__()
self.margin = margin
self.pdist = PairwiseDistance(2)
def forward(self, anchor, positive, negative):
# [在triplet_loss.py 的 13行位置处，填写代码，完成triplet loss的计算]
# 法一:
loss = torch.clamp(self.pdist.forward(anchor, positive)-self.pdist.forward(anchor, negative)+self.margin, 0.0)
loss = torch.mean(loss)
return loss
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 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
class TripletLoss(nn.Module):
def __init__(self):
super(TripletLoss, self).__init__()
self.pdist = PairwiseDistance(2)
def forward(self, anchor, positive, negative):
pos_dist = self.pdist.forward(anchor, positive)
neg_dist = self.pdist.forward(anchor, negative)
return torch.mean(logistic(neg_dist - pos_dist))
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()
class Recognition:
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
def detect(self, img1, img2, threshhold):
img1 = img1.cuda()
img2 = img2.cuda()
output1, output2 = self.model(img1), self.model(img2)
output1 = torch.div(output1, torch.norm(output1))
output2 = torch.div(output2, torch.norm(output2))
distance = self.l2_distance.forward(output1, output2)
if distance < threshhold:
self.is_same = True
return self.is_same
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
class ContrastiveLoss(torch.nn.Module):
def __init__(self, margin=1):
super(ContrastiveLoss, self).__init__()
self.margin = margin
self.pdist = PairwiseDistance(2)
def forward(self, output1, output2, label):
l2_dist = self.pdist.forward(output1, output2)
l2_dist = l2_dist.float()
label = label.float()
loss_contrastive = torch.mean(label * torch.pow(l2_dist, 2) +
(1-label) * torch.pow(torch.clamp(self.margin - l2_dist, min=0.0), 2))
return loss_contrastive
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
class TripletLoss(torch.autograd.Function):
def __init__(self, margin):
super(TripletLoss, self).__init__()
self.margin = margin
self.pdist = PairwiseDistance(2)
def forward(self, anchor, positive, negative):
"""
:param anchor: samples embedding from source domain. n x d: n is sample number, d is dimension
:param positive: hard positives
:param negative: hard negatives
:return:
"""
pos_dist = self.pdist.forward(anchor, positive)
neg_dist = self.pdist.forward(anchor, negative)
hinge_dist = torch.clamp(self.margin + pos_dist - neg_dist, min = 0.0) # Clamp all elements in input into the range [ 0.0, self.margin + pos_dist - neg_dist ] and return a resulting tensor
loss = torch.mean(hinge_dist)
return loss
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)
# 模型加载
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 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 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 __init__(self, margin):
super(TripletLoss, self).__init__()
self.margin = margin
self.pdist = PairwiseDistance(2)
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))
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[ ]:
for param_group in optimizer.param_groups: # 调整学习率
param_group['lr'] = lr
lr = 0.125
optimizer_model = torch.optim.Adagrad(model.parameters(),
lr=lr,
lr_decay=1e-4,
weight_decay=0)
# 打卡时间、epoch
start_epoch = 0
total_time_start = time.time()
end_epoch = start_epoch + 300
# 导入l2计算的
l2_distance = PairwiseDistance(2)
# 为了打日志先预制个最低auc和最佳acc在前头
best_roc_auc = -1
best_accuracy = -1
# epoch大循环
for epoch in range(start_epoch, end_epoch):
print("\ntraining on TrainDataset! ...")
epoch_time_start = time.time()
triplet_loss_sum = 0
sample_num = 0
model.train() # 训练模式
# step小循环
progress_bar = enumerate(tqdm(train_dataloader))
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 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))
