|
import torch |
|
from models import FBankCrossEntropyNet |
|
import tqdm |
|
import multiprocessing |
|
import time |
|
import numpy as np |
|
from models import DynamicLinearClassifier |
|
MODEL_PATH = './weights/triplet_loss_trained_model.pth' |
|
model_instance = FBankCrossEntropyNet() |
|
model_instance.load_state_dict(torch.load(MODEL_PATH, map_location=lambda storage, loc: storage)) |
|
|
|
use_cuda = False |
|
kwargs = {'num_workers': multiprocessing.cpu_count(), |
|
'pin_memory': True} if use_cuda else {} |
|
|
|
|
|
def train_classification(model, device, train_loader, optimizer, epoch, log_interval): |
|
model.train() |
|
losses = [] |
|
accuracy = 0 |
|
for batch_idx, (x, y) in enumerate(tqdm.tqdm(train_loader)): |
|
x, y = x.to(device), y.to(device) |
|
x = model_instance(x) |
|
optimizer.zero_grad() |
|
out = model(x) |
|
loss = model.loss(out, y) |
|
|
|
with torch.no_grad(): |
|
pred = torch.argmax(out, dim=1) |
|
accuracy += torch.sum((pred == y)) |
|
|
|
losses.append(loss.item()) |
|
loss.backward() |
|
optimizer.step() |
|
|
|
if batch_idx % log_interval == 0: |
|
print('{} Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( |
|
time.ctime(time.time()), |
|
epoch, batch_idx * len(x), len(train_loader.dataset), |
|
100. * batch_idx / len(train_loader), loss.item())) |
|
|
|
accuracy_mean = (100. * accuracy) / len(train_loader.dataset) |
|
|
|
return np.mean(losses), accuracy_mean.item() |
|
|
|
|
|
|
|
|
|
def test_classification(model, device, test_loader, log_interval=None): |
|
model.eval() |
|
losses = [] |
|
|
|
accuracy = 0 |
|
with torch.no_grad(): |
|
for batch_idx, (x, y) in enumerate(tqdm.tqdm(test_loader)): |
|
x, y = x.to(device), y.to(device) |
|
x = model_instance(x) |
|
out = model(x) |
|
test_loss_on = model.loss(out, y).item() |
|
losses.append(test_loss_on) |
|
|
|
pred = torch.argmax(out, dim=1) |
|
accuracy += torch.sum((pred == y)) |
|
|
|
if log_interval is not None and batch_idx % log_interval == 0: |
|
print('{} Test: [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( |
|
time.ctime(time.time()), |
|
batch_idx * len(x), len(test_loader.dataset), |
|
100. * batch_idx / len(test_loader), test_loss_on)) |
|
|
|
test_loss = np.mean(losses) |
|
accuracy_mean = (100. * accuracy) / len(test_loader.dataset) |
|
|
|
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} , ({:.4f})%\n'.format( |
|
test_loss, accuracy, len(test_loader.dataset), accuracy_mean)) |
|
return test_loss, accuracy_mean.item() |
|
|
|
|
|
|
|
def speaker_probability(tensor): |
|
counts = {} |
|
total = 0 |
|
for value in tensor: |
|
value = int(value) |
|
counts[value] = counts.get(value, 0) + 1 |
|
total += 1 |
|
|
|
probabilities = {} |
|
for key, value in counts.items(): |
|
probabilities['speaker '+str(key)] = value / total |
|
|
|
return probabilities |
|
|
|
|
|
|
|
def inference_speaker_classification( |
|
file_speaker, |
|
num_class=3, |
|
num_layers= 2, |
|
model_instance=model_instance, |
|
model_path='saved_models_cross_entropy_classification/0.pth' |
|
): |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
from utils.preprocessing import extract_fbanks |
|
fbanks = extract_fbanks(file_speaker) |
|
model = DynamicLinearClassifier(num_layers =num_layers ,output_size=num_class) |
|
cpkt = torch.load(model_path) |
|
model.load_state_dict(cpkt) |
|
model = model.double() |
|
model.to(device) |
|
model_instance = model_instance.double() |
|
model_instance.eval() |
|
model_instance.to(device) |
|
with torch.no_grad(): |
|
x = torch.from_numpy(fbanks) |
|
embedings = model_instance(x.to(device)) |
|
|
|
|
|
output = model(embedings) |
|
output = torch.argmax(output,dim=-1) |
|
speaker_pro = speaker_probability(output) |
|
print(speaker_pro) |
|
return speaker_pro |
|
|
|
|
