def main():
# Views the training images and displays the distance on anchor-negative and anchor-positive
test_display_triplet_distance = False
# print the experiment configuration
print('\nparsed options:\n{}\n'.format(vars(args)))
print('\nNumber of Classes:\n{}\n'.format(len(train_dir.classes)))
# instantiate model and initialize weightsNUM_FEATURES
# TODO(xin): IMPORTANT load num_classes from checkpoint
model = DeepSpeakerModel(embedding_size=args.embedding_size,
num_classes=len(train_dir.classes),
feature_dim=num_features,
frame_dim=c.NUM_FRAMES)
if args.cuda:
model.cuda()
from torchsummary import summary
summary(model, (1, c.NUM_FRAMES, c.NUM_FEATURES))
# # More detailed information on model
# print(model)
optimizer = create_optimizer(model, args.lr)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print('=> loading checkpoint {}'.format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
print('=> no checkpoint found at {}'.format(args.resume))
start = args.start_epoch
end = start + args.epochs
train_loader = torch.utils.data.DataLoader(train_dir, batch_size=args.batch_size, shuffle=False, **kwargs)
test_loader = torch.utils.data.DataLoader(test_dir, batch_size=args.test_batch_size, shuffle=False, **kwargs)
for epoch in range(start, end):
if args.test_only:
test(test_loader, model, epoch)
return
train(train_loader, model, optimizer, epoch)
test(test_loader, model, epoch)
def main():
# Views the training images and displays the distance on anchor-negative and anchor-positive
test_display_triplet_distance = False
# print the experiment configuration
print('\nparsed options:\n{}\n'.format(vars(args)))
print('\nNumber of Classes:\n{}\n'.format(len(train_dir.classes)))
# instantiate model and initialize weights
model = DeepSpeakerModel(embedding_size=args.embedding_size,
num_classes=len(train_dir.classes))
if args.cuda:
model.cuda()
optimizer = create_optimizer(model, args.lr)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print('=> loading checkpoint {}'.format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
print('=> no checkpoint found at {}'.format(args.resume))
start = args.start_epoch
#start = 0
end = start + args.epochs
train_loader = torch.utils.data.DataLoader(train_dir,
batch_size=args.batch_size,
shuffle=False,
**kwargs)
for epoch in range(start, end):
train(train_loader, model, optimizer, epoch)
#test(test_loader, model, epoch)
#break;
if test_display_triplet_distance:
display_triplet_distance(model, train_loader,
LOG_DIR + "/train_{}".format(epoch))
def main(libri_dir=c.DATASET_DIR):
print('Looking for fbank features [.npy] files in {}.'.format(libri_dir))
libri = data_catalog(libri_dir)
# filename speaker_id
# 0 audio/LibriSpeechSamples/train-clean-100-npy/1-100-0001.npy 1
# 1 audio/LibriSpeechSamples/train-clean-100-npy/1-100-0002.npy 1
unique_speakers = libri['speaker_id'].unique() # 251 speaker
transform=transforms.Compose([transforms.ToTensor()])
train_dir = stochastic_mini_batch(libri)
train_loader = DataLoader(train_dir, batch_size=c.BATCH_SIZE, shuffle=True)
model = DeepSpeakerModel(embedding_size=c.EMBEDDING_SIZE,num_classes=c.NUM_SPEAKERS)
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0)
epoch = 0
model.cuda()
summary(model, input_size=(1, 160, 64))
for epoch in range(100):
model.train()
for batch_idx, (data_a, data_p, data_n,label_a,label_p,label_n) in tqdm(enumerate(train_loader)):
data_a, data_p, data_n = data_a.type(torch.FloatTensor),data_p.type(torch.FloatTensor),data_n.type(torch.FloatTensor)
data_a, data_p, data_n = data_a.cuda(), data_p.cuda(), data_n.cuda()
data_a, data_p, data_n = Variable(data_a), Variable(data_p), Variable(data_n)
out_a, out_p, out_n = model(data_a), model(data_p), model(data_n)
triplet_loss = TripletMarginLoss(0.2).forward(out_a, out_p, out_n)
loss = triplet_loss
# compute gradient and update weights
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('selected_triplet_loss', triplet_loss.data)
print("epoch:",epoch)
torch.save(model.state_dict(),"checkpoint_{}.pt".format(epoch))