def main():
if args.resume:
if not os.path.isfile('./checkpoint/{}.pkl'.format(args.model)):
raise ValueError('no models saved....!!!!')
print 'resume from checkpoint....'
net = torch.load('./checkpoint/{}.pkl'.format(args.model))
else:
if args.model == 'vgg16':
net = vgg.VGG(args.model)
elif args.model == 'vgg19':
net = vgg.VGG(args.model)
elif args.model == 'resnet18':
net = resnet.ResNet18()
elif args.model == 'resnet34':
net = resnet.ResNet34()
elif args.model == 'resnet50':
net = resnet.ResNet50()
elif args.model == 'resnet101':
net = resnet.ResNet101()
elif args.model == 'resnet152':
net = resnet.ResNet152()
elif args.model == 'densenet121':
net = densenet.DenseNet121()
elif args.model == 'densenet161':
net = densenet.DenseNet161()
elif args.model == 'densenet169':
net = densenet.DenseNet169()
elif args.model == 'densenet201':
net = densenet.DenseNet201()
else:
raise ValueError('model not implemented...!!')
net.cuda(args.gpu)
net = nn.DataParallel(net, device_ids = range(torch.cuda.device_count()))
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optim = torch.optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=100, gamma=0.1)
for e in xrange(args.epoch):
train(e, net, criterion, optim, lr_scheduler)
test(e, net)
idx_to_class = {}
for key in testset.spk2label:
idx_to_class[str(testset.spk2label[key].squeeze().item())] = key
print(idx_to_class, '\n')
ckpt = torch.load(args.cp_path, map_location=lambda storage, loc: storage)
if args.model == 'cnn':
model = base_cnn.CNN(n_classes=args.nclasses)
elif args.model == 'vgg':
model = vgg.VGG('VGG11', n_classes=args.nclasses)
elif args.model == 'resnet':
model = resnet.ResNet12(n_classes=args.nclasses)
elif args.model == 'densenet':
model = densenet.DenseNet121(n_classes=args.nclasses)
elif args.model == 'tdnn':
model = TDNN.TDNN(n_classes=args.nclasses)
try:
print(model.load_state_dict(ckpt['model_state'], strict=True))
print('\n')
except RuntimeError as err:
print("Runtime Error: {0}".format(err))
except:
print("Unexpected error:", sys.exc_info()[0])
raise
print('\n\nNumber of parameters: {}\n'.format(
sum(p.numel() for p in model.parameters())))
std=[0.229, 0.224, 0.225])
])
validset = datasets.ImageFolder(args.data_path, transform=transform_test)
valid_loader = torch.utils.data.DataLoader(validset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_workers)
labels_list = [x[1] for x in validset]
if args.model == 'vgg':
model = vgg.VGG('VGG19')
elif args.model == 'resnet':
model = resnet.ResNet50()
elif args.model == 'DenseNet121':
model = densenet.DenseNet121()
ckpt = torch.load(args.cp_path, map_location=lambda storage, loc: storage)
try:
model.load_state_dict(ckpt['model_state'], strict=True)
except RuntimeError as err:
print("Runtime Error: {0}".format(err))
except:
print("Unexpected error:", sys.exc_info()[0])
raise
if args.cuda:
device = get_freer_gpu()
model = model.cuda(device)
else:
device = torch.device('cpu')
n_classes=n_classes,
emb_size=emb_size,
r_proj_size=rproj_size)
elif args.model == 'resnet':
model = resnet.ResNet50(nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=n_classes,
emb_size=emb_size,
r_proj_size=rproj_size)
elif args.model == 'densenet':
model = densenet.DenseNet121(nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=n_classes,
emb_size=emb_size,
r_proj_size=rproj_size)
print(model.load_state_dict(ckpt['model_state'], strict=False))
if args.cuda:
device = get_freer_gpu()
model = model.cuda(device)
else:
device = torch.device('cpu')
embeddings = []
labels = []
def train(lr, l2, momentum, smoothing, warmup, model, emb_size, n_hidden,
hidden_size, dropout_prob, epochs, batch_size, valid_batch_size,
n_workers, cuda, data_path, valid_data_path, hdf_path,
valid_hdf_path, checkpoint_path, softmax, pretrained,
pretrained_path, max_gnorm, stats, log_dir, eval_every, ablation):
args_dict = locals()
cp_name = get_cp_name(checkpoint_path)
if pretrained_path != 'none':
print('\nLoading pretrained model from: {}\n'.format(
args.pretrained_path))
ckpt = torch.load(pretrained_path,
map_location=lambda storage, loc: storage)
dropout_prob, n_hidden, hidden_size, emb_size = ckpt[
'dropout_prob'], ckpt['n_hidden'], ckpt['hidden_size'], ckpt[
'emb_size']
if 'r_proj_size' in ckpt:
rproj_size = ckpt['r_proj_size']
else:
rproj_size = -1
print('\nUsing pretrained config for discriminator. Ignoring args.')
args_dict['dropout_prob'], args_dict['n_hidden'], args_dict[
'hidden_size'], args_dict[
'emb_size'] = dropout_prob, n_hidden, hidden_size, emb_size
if log_dir != 'none':
writer = SummaryWriter(log_dir=os.path.join(log_dir, cp_name),
comment=model,
purge_step=0)
writer.add_hparams(hparam_dict=args_dict,
metric_dict={'best_eer': 0.0})
else:
writer = None
if stats == 'cars':
mean, std = [0.4461, 0.4329, 0.4345], [0.2888, 0.2873, 0.2946]
elif stats == 'cub':
mean, std = [0.4782, 0.4925, 0.4418], [0.2330, 0.2296, 0.2647]
elif stats == 'sop':
mean, std = [0.5603, 0.5155, 0.4796], [0.2939, 0.2991, 0.3085]
elif stats == 'imagenet':
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
if hdf_path != 'none':
transform_train = transforms.Compose([
transforms.ToPILImage(),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(10),
transforms.RandomPerspective(p=0.1),
transforms.RandomGrayscale(p=0.1),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
trainset = Loader(hdf_path, transform_train)
else:
transform_train = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(10),
transforms.RandomPerspective(p=0.1),
transforms.RandomGrayscale(p=0.1),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
trainset = datasets.ImageFolder(data_path, transform=transform_train)
train_loader = torch.utils.data.DataLoader(
trainset,
batch_size=batch_size,
shuffle=True,
num_workers=n_workers,
worker_init_fn=set_np_randomseed,
pin_memory=True)
if valid_hdf_path != 'none':
transform_test = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
validset = Loader(args.valid_hdf_path, transform_test)
else:
transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
validset = datasets.ImageFolder(args.valid_data_path,
transform=transform_test)
valid_loader = torch.utils.data.DataLoader(validset,
batch_size=valid_batch_size,
shuffle=True,
num_workers=n_workers,
pin_memory=True)
nclasses = trainset.n_classes if isinstance(trainset, Loader) else len(
trainset.classes)
if model == 'vgg':
model_ = vgg.VGG('VGG19',
nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=nclasses,
emb_size=emb_size,
r_proj_size=rproj_size)
elif model == 'resnet':
model_ = resnet.ResNet50(nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=nclasses,
emb_size=emb_size,
r_proj_size=rproj_size)
elif model == 'densenet':
model_ = densenet.DenseNet121(nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=nclasses,
emb_size=emb_size,
r_proj_size=rproj_size)
if pretrained_path != 'none':
if ckpt['sm_type'] == 'am_softmax':
del (ckpt['model_state']['out_proj.w'])
elif ckpt['sm_type'] == 'softmax':
del (ckpt['model_state']['out_proj.w.weight'])
del (ckpt['model_state']['out_proj.w.bias'])
print(model_.load_state_dict(ckpt['model_state'], strict=False))
print('\n')
if pretrained:
print('\nLoading pretrained encoder from torchvision\n')
if model == 'vgg':
model_pretrained = torchvision.models.vgg19(pretrained=True)
elif model == 'resnet':
model_pretrained = torchvision.models.resnet50(pretrained=True)
elif model == 'densenet':
model_pretrained = torchvision.models.densenet121(pretrained=True)
print(
model_.load_state_dict(model_pretrained.state_dict(),
strict=False))
print('\n')
if cuda:
device = get_freer_gpu()
model_ = model_.cuda(device)
torch.backends.cudnn.benchmark = True
optimizer = TransformerOptimizer(optim.SGD(model_.parameters(),
lr=lr,
momentum=momentum,
weight_decay=l2,
nesterov=True),
lr=lr,
warmup_steps=warmup)
trainer = TrainLoop(model_,
optimizer,
train_loader,
valid_loader,
max_gnorm=max_gnorm,
label_smoothing=smoothing,
verbose=-1,
cp_name=cp_name,
save_cp=True,
checkpoint_path=checkpoint_path,
ablation=ablation,
cuda=cuda,
logger=writer)
for i in range(5):
print(' ')
print('Hyperparameters:')
print('Selected model: {}'.format(model))
print('Embedding size: {}'.format(emb_size))
print('Hidden layer size: {}'.format(hidden_size))
print('Number of hidden layers: {}'.format(n_hidden))
print('Random projection size: {}'.format(rproj_size))
print('Dropout rate: {}'.format(dropout_prob))
print('Batch size: {}'.format(batch_size))
print('LR: {}'.format(lr))
print('Momentum: {}'.format(momentum))
print('l2: {}'.format(l2))
print('Label smoothing: {}'.format(smoothing))
print('Warmup iterations: {}'.format(warmup))
print('Softmax Mode is: {}'.format(softmax))
print('Pretrained: {}'.format(pretrained))
print('Pretrained path: {}'.format(pretrained_path))
print('Evaluate every {} iterations.'.format(eval_every))
print('Ablation Mode: {}'.format(ablation))
print(' ')
if i > 0:
print(' ')
print('Trial {}'.format(i + 1))
print(' ')
try:
cost = trainer.train(n_epochs=epochs,
save_every=epochs + 10,
eval_every=eval_every)
print(' ')
print('Best e2e EER in file ' + cp_name +
' was: {}'.format(cost[0]))
print('Best cos EER in file ' + cp_name +
' was: {}'.format(cost[1]))
print(' ')
if log_dir != 'none':
writer.add_hparams(hparam_dict=args_dict,
metric_dict={'best_eer': cost[0]})
return cost[0]
except:
print("Error:", sys.exc_info())
pass
print('Returning dummy cost due to failures while training.')
return 0.99
n_h=args.hidden_size,
dropout_prob=args.dropout_prob,
sm_type=args.softmax,
n_classes=args.nclasses,
emb_size=args.emb_size)
elif args.model == 'resnet':
model = resnet.ResNet50(nh=args.n_hidden,
n_h=args.hidden_size,
dropout_prob=args.dropout_prob,
sm_type=args.softmax,
n_classes=args.nclasses,
emb_size=args.emb_size)
elif args.model == 'densenet':
model = densenet.DenseNet121(nh=args.n_hidden,
n_h=args.hidden_size,
dropout_prob=args.dropout_prob,
sm_type=args.softmax,
n_classes=args.nclasses,
emb_size=args.emb_size)
if args.pretrained_path:
if ckpt['sm_type'] == 'am_softmax':
del (ckpt['model_state']['out_proj.w'])
elif ckpt['sm_type'] == 'softmax':
del (ckpt['model_state']['out_proj.w.weight'])
del (ckpt['model_state']['out_proj.w.bias'])
print(model.load_state_dict(ckpt['model_state'], strict=False))
print('\n')
elif args.pretrained:
print('\nLoading pretrained encoder from torchvision\n')
])
validset = datasets.ImageFolder(args.valid_data_path,
transform=transform_test)
valid_loader = torch.utils.data.DataLoader(validset,
batch_size=args.valid_batch_size,
shuffle=True,
num_workers=args.n_workers,
pin_memory=True)
if args.model == 'vgg':
model = vgg.VGG('VGG19', sm_type=args.softmax)
elif args.model == 'resnet':
model = resnet.ResNet50(sm_type=args.softmax)
elif args.model == 'densenet':
model = densenet.DenseNet121(sm_type=args.softmax)
if args.cuda:
device = get_freer_gpu()
model = model.to(device)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
weight_decay=args.l2,
momentum=args.momentum)
trainer = TrainLoop(model,
optimizer,
train_loader,
valid_loader,
margin=args.margin,
if args.valid_hdf_path:
transform_test = transforms.Compose([transforms.ToPILImage(), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
validset = Loader(args.valid_hdf_path, transform_train)
else:
transform_test = transforms.Compose([transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
validset = datasets.ImageFolder(args.valid_data_path, transform=transform_test)
valid_loader = torch.utils.data.DataLoader(validset, batch_size=args.valid_batch_size, shuffle=True, num_workers=args.n_workers, pin_memory=True)
if args.model == 'vgg':
model = vgg.VGG('VGG19', sm_type=args.softmax, n_classes=args.nclasses)
elif args.model == 'resnet':
model = resnet.ResNet50(sm_type=args.softmax, n_classes=args.nclasses)
elif args.model == 'densenet':
model = densenet.DenseNet121(sm_type=args.softmax, n_classes=args.nclasses)
if args.pretrained:
print('\nLoading pretrained model\n')
if args.model == 'vgg':
model_pretrained = torchvision.models.vgg19(pretrained=True)
elif args.model == 'resnet':
model_pretrained = torchvision.models.resnet50(pretrained=True)
elif args.model == 'densenet':
model_pretrained = torchvision.models.densenet121(pretrained=True)
print(model.load_state_dict(model_pretrained.state_dict(), strict=False))
print('\n')
if args.cuda:
device = get_freer_gpu()
def main():
# Training settings
parser = argparse.ArgumentParser()
parser.add_argument('--batch-size',
type=int,
default=64,
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size',
type=int,
default=64,
help='input batch size for testing (default: 64)')
parser.add_argument('--epochs',
type=int,
default=100,
help='number of epochs to train (default: 100)')
parser.add_argument('--lr',
type=float,
default=0.01,
help='learning rate (default: 0.01)')
parser.add_argument('--momentum',
type=float,
default=0.9,
help='SGD momentum (default: 0.9)')
parser.add_argument(
'--log-interval',
type=int,
default=5,
help=
'how many batches to wait before logging training status (default: 5)')
parser.add_argument('--seed',
type=int,
default=42,
help='random seed (default: 42)')
parser.add_argument('--use-checkpoint',
action='store_true',
default=False,
help='enables loading model from checkpoint')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device('cuda' if use_cuda else 'cpu')
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
start_epoch = 0
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=transform_train)
train_loader = torch.utils.data.DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
testset = datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform_test)
test_loader = torch.utils.data.DataLoader(testset,
batch_size=args.test_batch_size,
shuffle=False,
**kwargs)
model = densenet.DenseNet121().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum)
if not os.path.isdir('checkpoints'):
os.mkdir('checkpoints')
if args.use_checkpoint:
checkpoint = torch.load('./checkpoints/densenet121')
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
for epoch in range(start_epoch, args.epochs):
train(args, model, device, train_loader, criterion, optimizer, epoch)
test(args, model, device, test_loader, criterion, epoch)
def train(lr, l2, momentum, smoothing, patience, model, emb_size, n_hidden,
hidden_size, dropout_prob, epochs, batch_size, valid_batch_size,
n_workers, cuda, data_path, hdf_path, valid_data_path,
valid_hdf_path, checkpoint_path, softmax, n_classes, pretrained,
max_gnorm, lr_factor):
cp_name = get_cp_name(checkpoint_path)
if hdf_path != 'none':
transform_train = transforms.Compose([
transforms.ToPILImage(),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(30),
transforms.RandomPerspective(p=0.2),
transforms.ColorJitter(brightness=2),
transforms.RandomGrayscale(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
trainset = Loader(hdf_path, transform_train)
else:
transform_train = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(30),
transforms.RandomPerspective(p=0.2),
transforms.ColorJitter(brightness=2),
transforms.RandomGrayscale(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
trainset = datasets.ImageFolder(data_path, transform=transform_train)
train_loader = torch.utils.data.DataLoader(
trainset,
batch_size=batch_size,
shuffle=True,
num_workers=n_workers,
worker_init_fn=set_np_randomseed,
pin_memory=True)
if valid_hdf_path != 'none':
transform_test = transforms.Compose([
transforms.ToPILImage(),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
validset = Loader(args.valid_hdf_path, transform_test)
else:
transform_test = transforms.Compose([
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
validset = datasets.ImageFolder(args.valid_data_path,
transform=transform_test)
valid_loader = torch.utils.data.DataLoader(validset,
batch_size=valid_batch_size,
shuffle=True,
num_workers=n_workers,
pin_memory=True)
if model == 'vgg':
model_ = vgg.VGG('VGG19',
nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=n_classes,
emb_size=emb_size)
elif model == 'resnet':
model_ = resnet.ResNet50(nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=n_classes,
emb_size=emb_size)
elif model == 'densenet':
model_ = densenet.DenseNet121(nh=n_hidden,
n_h=hidden_size,
dropout_prob=dropout_prob,
sm_type=softmax,
n_classes=n_classes,
emb_size=emb_size)
if pretrained:
print('\nLoading pretrained encoder from torchvision\n')
if model == 'vgg':
model_pretrained = torchvision.models.vgg19(pretrained=True)
elif model == 'resnet':
model_pretrained = torchvision.models.resnet50(pretrained=True)
elif model == 'densenet':
model_pretrained = torchvision.models.densenet121(pretrained=True)
print(
model_.load_state_dict(model_pretrained.state_dict(),
strict=False))
print('\n')
if cuda:
device = get_freer_gpu()
model_ = model_.cuda(device)
torch.backends.cudnn.benchmark = True
optimizer = optim.SGD(model_.parameters(),
lr=lr,
weight_decay=l2,
momentum=momentum)
trainer = TrainLoop(model_,
optimizer,
train_loader,
valid_loader,
max_gnorm=max_gnorm,
patience=int(patience),
lr_factor=lr_factor,
label_smoothing=smoothing,
verbose=-1,
cp_name=cp_name,
save_cp=True,
checkpoint_path=checkpoint_path,
cuda=cuda)
for i in range(5):
print(' ')
print('Hyperparameters:')
print('Selected model: {}'.format(model))
print('Embedding size: {}'.format(emb_size))
print('Hidden layer size: {}'.format(hidden_size))
print('Number of hidden layers: {}'.format(n_hidden))
print('Dropout rate: {}'.format(dropout_prob))
print('Batch size: {}'.format(batch_size))
print('LR: {}'.format(lr))
print('Momentum: {}'.format(momentum))
print('l2: {}'.format(l2))
print('Label smoothing: {}'.format(smoothing))
print('Patience: {}'.format(patience))
print('Softmax Mode is: {}'.format(softmax))
print('Pretrained: {}'.format(pretrained))
print(' ')
if i > 0:
print(' ')
print('Trial {}'.format(i + 1))
print(' ')
try:
cost = trainer.train(n_epochs=epochs, save_every=epochs + 10)
print(' ')
print('Best e2e EER in file ' + cp_name +
' was: {}'.format(cost[0]))
print('Best cos EER in file ' + cp_name +
' was: {}'.format(cost[1]))
print(' ')
return cost[0]
except:
print("Error:", sys.exc_info())
pass
print('Returning dummy cost due to failures while training.')
return 0.99
metavar='N',
help='Embedding dimension (default: 256)')
args = parser.parse_args()
if args.model == 'vgg':
model = vgg.VGG('VGG19',
nh=args.n_hidden,
n_h=args.hidden_size,
emb_size=args.emb_size)
elif args.model == 'resnet':
model = resnet.ResNet50(nh=args.n_hidden,
n_h=args.hidden_size,
emb_size=args.emb_size)
elif args.model == 'densenet':
model = densenet.DenseNet121(nh=args.n_hidden,
n_h=args.hidden_size,
emb_size=args.emb_size)
batch = torch.rand(3, 3, 224, 224)
emb, out = model.forward(batch)
print(emb.size(), out.size())
out_layer = model.out_proj(out)
print(out_layer.size())
emb = torch.cat([emb, emb], 1)
print(emb.size())
