def get_traindataloaders(dataset_name: str,
config) -> data.DataLoader:
transfrom_list = [transforms.Resize((config.hyperparameters.image_size, config.hyperparameters.image_size), interpolation=1)]
if config.hyperparameters.random_hor_flip:
transfrom_list.append(transforms.RandomHorizontalFlip())
transfrom_list.append(transforms.ToTensor())
train_transforms = transforms.Compose(transfrom_list)
fold_tool = utils.FoldGenerator(config.dataset.cross_validation.num_fold,
config.dataset.cross_validation.num_train_folds,
config.dataset.cross_validation.num_val_folds)
train_folds, val_folds, test_folds = fold_tool.get_fold()
# VGGFACE2 dataset
if dataset_name == 'vggface2':
data_loader = dataset_utils.vggface2.get_vggface2_trainset(config.dataset.vggface2.train_dir,
train_transforms,
config.hyperparameters.people_per_batch,
config.hyperparameters.images_per_person)
# COXS2V dataset
elif 'cox_video' in dataset_name:
data_loader = dataset_utils.coxs2v.get_coxs2v_trainset(dataset_name,
config.dataset.coxs2v,
train_folds,
config.dataset.cross_validation.num_fold,
train_transforms,
config.hyperparameters.people_per_batch,
config.hyperparameters.images_per_person)
# BBT generated dataset
elif dataset_name == 'bbt':
data_loader = dataset_utils.bbt.get_trainset(config.dataset.bbt.dataset_path,
config.hyperparameters.images_per_person,
transform=train_transforms)
# Movie generated dataset
elif dataset_name == 'movie':
data_loader = dataset_utils.movie.get_trainset(config.dataset.movie.dataset_path,
config.hyperparameters.images_per_person,
transform=train_transforms)
# MNIST dataset
elif dataset_name == "mnist":
data_loader = dataset_utils.mnist.get_mnist_trainset(batch_size=config.hyperparameters.batch_size,
small=True)
elif dataset_name == "mnistbig":
data_loader = dataset_utils.mnist.get_mnist_trainset(batch_size=config.hyperparameters.batch_size,
small=False)
# USPS dataset
elif dataset_name == 'usps':
data_loader = dataset_utils.usps.get_usps_trainset(batch_size=config.hyperparameters.batch_size,
small=True)
elif dataset_name == 'uspsbig':
data_loader = dataset_utils.usps.get_usps_trainset(batch_size=config.hyperparameters.batch_size,
small=False)
else:
raise Exception('Datatset {} not supported.'.format(dataset_name))
return data_loader
def get_evaluators(dataset_list: list,
config) -> list:
evaluators_list = []
eval_transforms = transforms.Compose([
transforms.Resize((config.hyperparameters.image_size, config.hyperparameters.image_size), interpolation=1),
transforms.ToTensor()
])
fold_tool = utils.FoldGenerator(config.dataset.cross_validation.num_fold,
config.dataset.cross_validation.num_train_folds,
config.dataset.cross_validation.num_val_folds)
train_folds, val_folds, test_folds = fold_tool.get_fold()
for dataset_name in dataset_list:
# VGGFACE2 dataset
if 'vggface2' == dataset_name:
evaluator = dataset_utils.vggface2.get_vggface2_evaluator(config.dataset.vggface2.test_dir,
config.dataset.vggface2.pairs_file,
eval_transforms,
config.hyperparameters.batch_size,
preload=False)
# LFW dataset
elif 'lfw' == dataset_name:
evaluator = dataset_utils.lfw.get_evaluator(config.dataset.lfw,
eval_transforms,
config.hyperparameters.batch_size,
preload=False)
# COXS2V dataset
elif 'cox_video' in dataset_name:
evaluator = dataset_utils.coxs2v.get_evaluator(dataset_name,
config.dataset.coxs2v,
test_folds,
config.dataset.cross_validation.num_fold,
eval_transforms,
config.hyperparameters.batch_size,
preload=True)
# BBT generated dataset
elif 'bbt' == dataset_name:
from dataset_utils import bbt
evaluator = bbt.get_evaluator(config.dataset.bbt.dataset_path,
eval_transforms,
config.hyperparameters.batch_size)
# MNIST dataset
elif 'mnist' == dataset_name:
evaluator = dataset_utils.mnist.get_evaluator(batch_size=config.hyperparameters.batch_size)
# USPS dataset
elif 'usps' == dataset_name:
evaluator = dataset_utils.usps.get_evaluator(batch_size=config.hyperparameters.batch_size)
else:
raise Exception('Datatset {} not supported.'.format(dataset_name))
evaluators_list.append(evaluator)
return evaluators_list
checkpoint_path = config.model.checkpoint_path
else:
checkpoint_path = None
model = models.load_model(
config.model.model_arch,
device,
checkpoint_path=checkpoint_path,
embedding_size=config.model.embedding_size,
imgnet_pretrained=config.model.pretrained_imagenet)
batch_size = (config.hyperparameters.people_per_batch *
config.hyperparameters.images_per_person) // 2
nrof_folds = config.dataset.cross_validation.num_fold
fold_tool = utils.FoldGenerator(
nrof_folds, config.dataset.cross_validation.num_train_folds,
config.dataset.cross_validation.num_val_folds)
train_folds, val_folds, test_folds = fold_tool.get_fold()
#Data transform
data_transform = transforms.Compose([
transforms.Resize((config.hyperparameters.image_size,
config.hyperparameters.image_size),
interpolation=1),
transforms.ToTensor()
])
# Get data loaders
test_loaders_list = dataset_utils.dataloaders.get_testevaluators(
config, data_transform, batch_size, test_folds, nrof_folds,
['bbt_ep01'])
import pandas as pd
import misc as ms
import test
import torch
import torch.nn as nn
import losses
from sklearn.cluster import KMeans
from torchvision import transforms
from utils import visdom_util
from dataset_utils import get_coxs2v_testset, get_coxs2v_trainset
import utils
fold_tool = utils.FoldGenerator(10, 2, 1)
train_folds, val_folds, test_folds = fold_tool.get_fold()
train_folds = [3, 4, 5, 6, 7, 8, 9]
test_folds = [3, 4, 5, 6, 7, 8, 9]
eval_folds = [3, 4, 5, 6, 7, 8, 9]
people_per_batch = 20
images_per_person = 5
def compare_models(model_1, model_2):
models_differ = 0
for key_item_1, key_item_2 in zip(model_1.state_dict().items(),
model_2.state_dict().items()):
if torch.equal(key_item_1[1], key_item_2[1]):
pass
else:
models_differ += 1
def main(args):
print('CONFIGURATION:\t{}'.format(args.config))
with open(args.config) as json_config_file:
config = utils.AttrDict(json.load(json_config_file))
# Set up output directory
# subdir = datetime.strftime(datetime.now(), '%Y%m%d-%H%M%S')
subdir = config.visdom.environment_name
model_dir = os.path.join(os.path.expanduser(config.output.output_dir), subdir)
if not os.path.isdir(model_dir): # Create the model directory if it doesn't exist
os.makedirs(model_dir)
else:
raise Exception('Environment name {} already taken.'.format(subdir))
config_filename = path_leaf(args.config)
copyfile(args.config, os.path.join(model_dir, config_filename))
# CUDA for PyTorch
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
data_transform = transforms.Compose([
transforms.Resize((config.hyperparameters.image_size, config.hyperparameters.image_size), interpolation=1),
transforms.ToTensor()
])
test_batch_size = (config.hyperparameters.people_per_batch * config.hyperparameters.images_per_person) // 2
nrof_folds = config.dataset.cross_validation.num_fold
fold_tool = utils.FoldGenerator(nrof_folds,
config.dataset.cross_validation.num_train_folds,
config.dataset.cross_validation.num_val_folds)
train_folds, val_folds, test_folds = fold_tool.get_fold()
###########################
# SET UP DATALOADERS HERE #
###########################
source_loader = dataset_utils.get_coxs2v_trainset(config.dataset.coxs2v.still_dir,
config.dataset.coxs2v.video2_dir,
config.dataset.coxs2v.video2_pairs,
train_folds,
nrof_folds,
data_transform,
config.hyperparameters.people_per_batch,
config.hyperparameters.images_per_person)
target_loader = dataset_utils.get_coxs2v_trainset(config.dataset.coxs2v.still_dir,
config.dataset.coxs2v.video4_dir,
config.dataset.coxs2v.video4_pairs,
val_folds,
nrof_folds,
data_transform,
config.hyperparameters.people_per_batch,
config.hyperparameters.images_per_person)
test_loaders_list = dataloaders.get_testevaluators(config,
data_transform,
test_batch_size,
test_folds,
nrof_folds,
is_vggface2=False,
is_lfw=True,
is_cox_video1=False,
is_cox_video2=True,
is_cox_video3=False,
is_cox_video4=True)
###################
# DATALOADERS END #
###################
#Set up training model
print('Building training model')
if config.model.checkpoint:
print('Loading from checkpoint {}'.format(config.model.checkpoint_path))
checkpoint = torch.load(config.model.checkpoint_path)
embedding_size = checkpoint['embedding_size']
# start_epoch = checkpoint['epoch']
start_epoch = 0
else:
embedding_size = config.model.embedding_size
start_epoch = 0
model = models.load_model(config.model.model_arch,
embedding_size=embedding_size,
imgnet_pretrained=config.model.pretrained_imagenet)
optimizer = optim.SGD(model.parameters(), lr=config.hyperparameters.learning_rate, momentum=0.9, nesterov=True, weight_decay=2e-4)
scheduler = lr_scheduler.ExponentialLR(optimizer, config.hyperparameters.learning_rate_decay_factor)
if config.model.checkpoint:
model.load_state_dict(checkpoint['model_state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
# scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
model = model.to(device)
plotter = utils.VisdomPlotter(env_name=config.visdom.environment_name, port=config.visdom.port)
print('Quadruplet loss training mode.')
miner = miners.SemihardNegativeQuadrupletSelector(config.hyperparameters.margin)
loss = losses.QuadrupletLoss2(config.hyperparameters.margin,
config.hyperparameters.margin2,
lamda=config.hyperparameters.lamda)
trainer = Quadruplet_Trainer(model,
miner,
loss,
optimizer,
scheduler,
device,
plotter,
config.hyperparameters.margin,
config.model.embedding_size,
config.visdom.log_interval)
# Loop over epochs
print('Training Launched.')
for epoch in range(start_epoch, config.hyperparameters.n_epochs):
# Validation
for test_name, test_loader in test_loaders_list:
print('\nEvaluation on {}'.format(test_name))
trainer.Evaluate(test_loader,
name=test_name,
nrof_folds=nrof_folds,
val_far=config.hyperparameters.val_far)
# Training
print('\nTrain Epoch {}'.format(epoch))
trainer.Train_Epoch(source_loader, target_loader)
# Save model
if not (epoch + 1) % config.output.save_interval:
if not os.path.isdir(model_dir): # Create the model directory if it doesn't exist
os.makedirs(model_dir)
model_file_path = os.path.join(model_dir, 'model_{}.pth'.format(epoch))
print('\nSave model at {}'.format(model_file_path))
torch.save({'epoch': epoch,
'model_state_dict': utils.state_dict_to_cpu(model.state_dict()),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'embedding_size': config.model.embedding_size
}, model_file_path)
print('Finish.')