def load_model(use_cuda, log_dir, cp_num, embedding_size, n_classes):
model = background_resnet(embedding_size=embedding_size, num_classes=n_classes)
if use_cuda:
model.cuda()
print('=> loading checkpoint')
# original saved file with DataParallel
checkpoint = torch.load(log_dir + '/checkpoint_' + str(cp_num) + '.pth')
# create new OrderedDict that does not contain `module.`
model.load_state_dict(checkpoint['state_dict'])
model.eval()
return model
def load_model(use_cuda, log_dir, cp_num, n_classes):
model = background_resnet(num_classes=n_classes)
if use_cuda:
model.cuda()
print('=> loading checkpoint')
# load pre-trained parameters
checkpoint = torch.load(log_dir + '/checkpoint_' + str(cp_num).zfill(3) + '.pth')
model.load_state_dict(checkpoint['state_dict'])
model.eval()
return model
def load_model(log_dir, cp_num, n_classes):
model = background_resnet(num_classes=n_classes)
optimizer = create_optimizer(model)
print('=> loading checkpoint')
checkpoint = torch.load(log_dir + '/checkpoint_' + str(cp_num).zfill(3) + '.pth')
# create new OrderedDict that does not contain `module.`
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
return model, optimizer
def main():
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
np.random.seed(args.seed)
cudnn.deterministic = True
if args.use_cuda:
torch.cuda.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
if args.use_checkpoint: start = args.cp_num + 1
else: start = 0 # Start epoch
n_epochs = max_epoch - start # How many epochs?
# Load dataset
train_DB, n_data, n_classes = make_DB(DB_type=args.data_type)
n_episode = int(n_data / ((args.n_shot + args.n_query) * args.nb_class_train))
if not os.path.exists(log_dir):
os.makedirs(log_dir)
# Generate model and optimizer
if args.use_checkpoint:
model, optimizer = load_model(log_dir, args.cp_num, n_classes)
else:
model = background_resnet(num_classes=n_classes)
optimizer = create_optimizer(model)
# define objective function, optimizer and scheduler
objective = Prototypical() if args.loss_type == 'prototypical' else SoftmaxLoss()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3, min_lr=1e-5, threshold=1e-4, verbose=1)
if args.use_cuda:
model.cuda()
train_generator = metaGenerator(train_DB, read_MFB,
nb_classes=args.nb_class_train, nb_samples_per_class=args.n_shot + args.n_query,
max_iter=n_episode * (n_epochs-args.cp_num), xp=np)
# training
train(train_generator, model, objective, optimizer, n_episode, log_dir, scheduler)
def main():
# Set hyperparameters
use_cuda = True # use gpu or cpu
val_ratio = 10 # Percentage of validation set
embedding_size = 128
start = 1 # Start epoch
n_epochs = 30 # How many epochs?
end = start + n_epochs # Last epoch
lr = 1e-1 # Initial learning rate
wd = 1e-4 # Weight decay (L2 penalty)
optimizer_type = 'sgd' # ex) sgd, adam, adagrad
batch_size = 64 # Batch size for training
valid_batch_size = 16 # Batch size for validation
use_shuffle = True # Shuffle for training or not
# Load dataset
train_dataset, valid_dataset, n_classes = load_dataset(val_ratio)
# print the experiment configuration
print('\nNumber of classes (speakers):\n{}\n'.format(n_classes))
log_dir = 'model_saved' # where to save checkpoints
if not os.path.exists(log_dir):
os.makedirs(log_dir)
# instantiate model and initialize weights
model = background_resnet(embedding_size=embedding_size, num_classes=n_classes)
if use_cuda:
model.cuda()
# define loss function (criterion), optimizer and scheduler
criterion = nn.CrossEntropyLoss()
optimizer = create_optimizer(optimizer_type, model, lr, wd)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, min_lr=1e-4, verbose=1)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=use_shuffle)
valid_loader = torch.utils.data.DataLoader(dataset=valid_dataset,
batch_size=valid_batch_size,
shuffle=False,
collate_fn = collate_fn_feat_padded)
# to track the average training loss per epoch as the model trains
avg_train_losses = []
# to track the average validation loss per epoch as the model trains
avg_valid_losses = []
for epoch in range(start, end):
# train for one epoch
train_loss = train(train_loader, model, criterion, optimizer, use_cuda, epoch, n_classes)
# evaluate on validation set
valid_loss = validate(valid_loader, model, criterion, use_cuda, epoch)
scheduler.step(valid_loss, epoch)
# calculate average loss over an epoch
avg_train_losses.append(train_loss)
avg_valid_losses.append(valid_loss)
# do checkpointing
torch.save({'epoch': epoch + 1, 'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()},
'{}/checkpoint_{}.pth'.format(log_dir, epoch))
# find position of lowest validation loss
minposs = avg_valid_losses.index(min(avg_valid_losses))+1
print('Lowest validation loss at epoch %d' %minposs)
# visualize the loss and learning rate as the network trained
visualize_the_losses(avg_train_losses, avg_valid_losses)
def main():
# GPU parmas
idx_cuda = 2
if idx_cuda < 0:
device = torch.device('cpu')
else:
device = torch.device(f'cuda:{idx_cuda}')
# Set hyperparameters
embedding_size = 128 # origial 128
start = 1 # Start epoch
n_epochs = 60 # How many epochs?
end = start + n_epochs # Last epoch
lr = 1e-1 # Initial learning rate
wd = 1e-4 # Weight decay (L2 penalty)
optimizer_type = 'adam' # ex) sgd, adam, adagrad
# The effective batch size is M*N
M = 5 # number of utterances per speaker
N = 10 # Number of speaker
print(f'Number of utternaces per speaker: {M}')
print(f'Number of speakers: {N}')
print(f'** Total batch size: {M*N} **')
use_shuffle = True # Shuffle for training or not
save_interval = 5
# Load dataset
train_dataset, n_classes = load_dataset(M)
log_dir = 'model_saved_verification_2' # where to save checkpoints
if not os.path.exists(log_dir):
os.makedirs(log_dir)
# instantiate model and initialize weights
model = background_resnet(embedding_size=embedding_size,
num_classes=n_classes,
backbone='resnet34')
# Load the wights trained for identification
#smodel.load_state_dict(torch.load('model_saved/checkpoint_50.pth')['state_dict'])
# remove the last layers
model.to(device)
# define loss function (criterion), optimizer and scheduler
criterion = GE2ELoss(device)
optimizer = create_optimizer(optimizer_type, model, lr, wd)
# Define dataloader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=N,
shuffle=use_shuffle,
collate_fn=my_collate)
# to track the average training loss per epoch as the model trains
avg_train_losses = []
for epoch in range(start, end):
# train for one epoch
train_loss = train(train_loader, model, criterion, optimizer, device,
epoch, N, M)
# calculate average loss over an epoch
avg_train_losses.append(train_loss)
# do checkpointing
if epoch % save_interval == 0 or epoch == end - 1:
torch.save(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, '{}/checkpoint_{}.pth'.format(log_dir, epoch))
