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))
import tensorflow as tf
import tqdm
if (__name__ == '__main__'):
torch.manual_seed(55)
torch.cuda.manual_seed_all(55)
if (__name__ == '__main__'):
model = DeepSpeakerModel(embedding_size=opt.embedding_size,
num_classes=opt.classes).cuda()
writer = SummaryWriter()
if (hasattr(opt, 'weights')):
pretrained_dict = torch.load(opt.weights)
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items()
if k in model_dict.keys() and v.size() == model_dict[k].size()
}
missed_params = [
k for k, v in model_dict.items()
if not k in pretrained_dict.keys()
]
print('loaded params/tot params:{}/{}'.format(len(pretrained_dict),
len(model_dict)))
print('miss matched params:{}'.format(missed_params))
model_dict.update(pretrained_dict)