| from trainer.fbankcross_classification import train_classification, test_classification, inference_speaker_classification |
| from utils.pt_util import restore_objects, save_model, save_objects, restore_model |
| import torch |
| from data_proc.cross_entropy_dataset import FBanksCrossEntropyDataset, DataLoader |
| import json |
| from torch import optim |
| import os |
| os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' |
| from models.classifier import DynamicLinearClassifier |
|
|
|
|
| async def train_csf( |
| train_dataset_path: str = 'dataset-speaker-csf/fbanks-train', |
| test_dataset_path: str = 'dataset-speaker-csf/fbanks-test', |
| model_name: str = 'fbanks-net-classification', |
| num_layers : int = 2 , |
| epoch: int = 2, |
| lr: float = 0.0005, |
| batch_size: int = 2, |
| labId: str = '', |
| ): |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| import multiprocessing |
| kwargs = {'num_workers': multiprocessing.cpu_count(), |
| 'pin_memory': True} if torch.cuda.is_available() else {} |
| try: |
| train_dataset = FBanksCrossEntropyDataset(train_dataset_path) |
| train_loader = DataLoader( |
| train_dataset, batch_size=batch_size, shuffle=True, **kwargs) |
| test_dataset = FBanksCrossEntropyDataset(test_dataset_path) |
| test_loader = DataLoader( |
| test_dataset, batch_size=batch_size, shuffle=True, **kwargs) |
| except: |
| return 'path dataset test or train is not exist' |
|
|
| try: |
|
|
| assert train_dataset.num_classes == test_dataset.num_classes |
|
|
| except: |
| return "The number of speakers in test and training sets must be equal " |
| if model_name == 'fbanks-net-classification': |
| try: |
| model = DynamicLinearClassifier(num_layers= num_layers, |
| output_size=train_dataset.num_classes).to(device) |
| except: |
| print('cuda load is error') |
| device = torch.device("cpu") |
| model = DynamicLinearClassifier(num_layers = num_layers, |
| output_size=train_dataset.num_classes).to(device) |
| else: |
| model = None |
| return {"model not exist in lab"} |
| model_path = f'./modelDir/{labId}/log_train/{model_name}/{num_layers}' |
| model = restore_model(model, model_path) |
| last_epoch, max_accuracy, train_losses, test_losses, train_accuracies, test_accuracies = restore_objects( |
| model_path, (0, 0, [], [], [], [])) |
| start = last_epoch + 1 if max_accuracy > 0 else 0 |
|
|
| models_path = [] |
| optimizer = optim.Adam(model.parameters(), lr) |
| for epoch in range(start, epoch): |
| train_loss, train_accuracy = train_classification( |
| model, device, train_loader, optimizer, epoch, 500) |
| test_loss, test_accuracy = test_classification( |
| model, device, test_loader) |
| print('After epoch: {}, train_loss: {}, test loss is: {}, train_accuracy: {}, ' |
| 'test_accuracy: {}'.format(epoch, train_loss, test_loss, train_accuracy, test_accuracy)) |
|
|
| train_losses.append(train_loss) |
| test_losses.append(test_loss) |
| train_accuracies.append(train_accuracy) |
| test_accuracies.append(test_accuracy) |
| if test_accuracy > max_accuracy: |
| max_accuracy = test_accuracy |
| model_path = save_model(model, epoch, model_path) |
| models_path.append(model_path) |
| save_objects((epoch, max_accuracy, train_losses, test_losses, |
| train_accuracies, test_accuracies), epoch, model_path) |
| print('saved epoch: {} as checkpoint'.format(epoch)) |
| train_history = { |
| "train_accuracies": train_accuracies, |
| "test_accuracies": test_accuracies, |
| "train_losses": train_losses, |
| "test_losses": test_losses, |
| "model_path": models_path |
| } |
| return { |
| 'history': json.dumps(train_history) |
| } |
|
|
|
|
| async def test_csf( |
| test_dataset_path: str = 'dataset-speaker-csf/fbanks-test', |
| model_name: str = 'fbanks-net-classification', |
| num_layers : int = 2, |
| batch_size: int = 2, |
| labId: str = '', |
| ): |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| import multiprocessing |
| kwargs = {'num_workers': multiprocessing.cpu_count(), |
| 'pin_memory': True} if torch.cuda.is_available() else {} |
| try: |
| test_dataset = FBanksCrossEntropyDataset(test_dataset_path) |
| test_loader = DataLoader( |
| test_dataset, batch_size=batch_size, shuffle=True, **kwargs) |
| except: |
| return 'path dataset test is not exist' |
| model_folder_path = f'./modelDir/{labId}/log_train/{model_name}/{num_layers}/' |
| for file in os.listdir(model_folder_path): |
| if file.endswith(".pth"): |
| model_path = os.path.join(model_folder_path, file) |
| if model_name == 'fbanks-net-classification': |
| try: |
| model = DynamicLinearClassifier(num_layers=num_layers, output_size=test_dataset.num_classes) |
| cpkt = torch.load(model_path) |
| model.load_state_dict(cpkt) |
| model.to(device) |
| except: |
| print('cuda load is error') |
| device = torch.device("cpu") |
| model = DynamicLinearClassifier(num_layers=num_layers,output_size=test_dataset.num_classes) |
| cpkt = torch.load(model_path) |
| model.load_state_dict(cpkt) |
| model.to(device) |
| else: |
| model = None |
| return {"model not exist in lab"} |
| test_loss, accurancy_mean = test_classification(model, device, test_loader) |
| print(accurancy_mean) |
| return { |
| 'test_loss': test_loss, |
| 'test_accuracy': accurancy_mean |
| } |
|
|
|
|
| def infer_csf( |
| speech_file_path: str = './sample.wav', |
| model_name: str = 'fbanks-net-classification', |
| num_layers : int = 2, |
| |
| labId: str = '', |
| ): |
| model_folder_path = f'./modelDir/{labId}/log_train/{model_name}/' |
| for file in os.listdir(model_folder_path): |
| if file.endswith(".pth"): |
| model_path = os.path.join(model_folder_path, file) |
| rs = inference_speaker_classification( |
| file_speaker=speech_file_path, model_path=model_path, num_layers = num_layers) |
| return { |
| "result": rs |
| } |
| |
| if __name__ == '__main__': |
| result = infer_csf() |
| print(result) |
