debug

nnet/xvector.py

@@ -55,6 +55,8 @@ from .loss import ArcLinear
 from .loss import ArcFace
 from .loss import l2_norm
 
+from torch.nn.parallel import DistributedDataParallel as DDP
+
 import tqdm
 
 __license__ = "LGPL"
@@ -442,7 +444,7 @@ class Xtractor(torch.nn.Module):
 
                 elif k.startswith("ctrans"):
                     segmental_layers.append((k, torch.nn.ConvTranspose1d(input_size,
-                                                                         cfg["segmental"][k]["output_channels"],
+                                                                         cfg["segmental"][k][":"],
                                                                          kernel_size=cfg["segmental"][k]["kernel_size"],
                                                                          dilation=cfg["segmental"][k]["dilation"])))
                 elif k.startswith("activation"):
@@ -905,7 +907,8 @@ def train_epoch(model, epoch, training_loader, optimizer, log_interval, device,
             accuracy += (torch.argmax(output.data, 1) == target).sum()
             if batch_idx % log_interval == 0:
                 batch_size = target.shape[0]
-                logging.critical('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tAccuracy: {:.3f}'.format(
+                logging.critical('{}, Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tAccuracy: {:.3f}'.format(
+                    time.strftime('%H:%M:%S', time.localtime()),
                     epoch, batch_idx + 1, training_loader.__len__(),
                     100. * batch_idx / training_loader.__len__(), loss.item(),
                     100.0 * accuracy.item() / ((batch_idx + 1) * batch_size)))
@@ -1101,3 +1104,286 @@ def extract_sliding_embedding(idmap_name,
                                  transform_pipeline=transform_pipeline)
 
     return embeddings
+
+
+def xdebug(speaker_number,
+           dataset_yaml,
+           epochs=None,
+           lr=None,
+           model_yaml=None,
+           model_name=None,
+           loss=None,
+           aam_margin=None,
+           aam_s=None,
+           patience=None,
+           tmp_model_name=None,
+           best_model_name=None,
+           multi_gpu=True,
+           clipping=False,
+           opt=None,
+           reset_parts=[],
+           freeze_parts=[],
+           num_thread=None):
+    """
+
+    :param speaker_number:
+    :param dataset_yaml:
+    :param epochs:
+    :param lr:
+    :param model_yaml:
+    :param model_name:
+    :param loss:
+    :param aam_margin:
+    :param aam_s:
+    :param patience:
+    :param tmp_model_name:
+    :param best_model_name:
+    :param multi_gpu:
+    :param clipping:
+    :param opt:
+    :param reset_parts:
+    :param freeze_parts:
+    :param num_thread:
+    :return:
+    """
+    # Add for tensorboard
+    # Writer will output to ./runs/ directory by default
+    # writer = SummaryWriter("runs/xvectors_experiments_2")
+    writer = None
+
+    if num_thread is None:
+        num_thread = multiprocessing.cpu_count()
+
+    logging.critical(f"Start process at {time.strftime('%H:%M:%S', time.localtime())}")
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Start from scratch
+    if model_name is None and model_yaml in ["xvector", "rawnet2"]:
+        # Initialize a first model
+        if model_yaml == "xvector":
+            model = Xtractor(speaker_number, "xvector")
+        elif model_yaml == "rawnet2":
+            model = Xtractor(speaker_number, "rawnet2")
+    else:
+        with open(model_yaml, 'r') as fh:
+            model_archi = yaml.load(fh, Loader=yaml.FullLoader)
+            if epochs is None:
+                epochs = model_archi["training"]["epochs"]
+            if patience is None:
+                patience = model_archi["training"]["patience"]
+            if opt is None:
+                opt = model_archi["training"]["opt"]
+            if lr is None:
+                lr = model_archi["training"]["lr"]
+            if loss is None:
+                loss = model_archi["training"]["loss"]
+            if aam_margin is None and model_archi["training"]["loss"] == "aam":
+                aam_margin = model_archi["training"]["aam_margin"]
+            if aam_s is None and model_archi["training"]["loss"] == "aam":
+                aam_s = model_archi["training"]["aam_s"]
+            if tmp_model_name is None:
+                tmp_model_name = model_archi["training"]["tmp_model_name"]
+            if best_model_name is None:
+                best_model_name = model_archi["training"]["best_model_name"]
+            if multi_gpu is None:
+                multi_gpu = model_archi["training"]["multi_gpu"]
+            if clipping is None:
+                clipping = model_archi["training"]["clipping"]
+
+        if model_name is None:
+            model = Xtractor(speaker_number, model_yaml)
+
+        # If we start from an existing model
+        else:
+            # Load the model
+            logging.critical(f"*** Load model from = {model_name}")
+            checkpoint = torch.load(model_name)
+            model = Xtractor(speaker_number, model_yaml)
+
+            """
+            Here we remove all layers that we don't want to reload
+
+            """
+            pretrained_dict = checkpoint["model_state_dict"]
+            for part in reset_parts:
+                pretrained_dict = {k: v for k, v in pretrained_dict.items() if not k.startswith(part)}
+
+            new_model_dict = model.state_dict()
+            new_model_dict.update(pretrained_dict)
+            model.load_state_dict(new_model_dict)
+
+        # Freeze required layers
+        for name, param in model.named_parameters():
+            if name.split(".")[0] in freeze_parts:
+                param.requires_grad = False
+
+    print(model)
+
+    if torch.cuda.device_count() > 1 and multi_gpu:
+        print("Let's use", torch.cuda.device_count(), "GPUs!")
+        model = torch.nn.DataParallel(model)
+    else:
+        print("Train on a single GPU")
+    model.to(device)
+
+    """
+    Set the dataloaders according to the dataset_yaml
+
+    First we load the dataframe from CSV file in order to split it for training and validation purpose
+    Then we provide those two 
+    """
+    with open(dataset_yaml, "r") as fh:
+        dataset_params = yaml.load(fh, Loader=yaml.FullLoader)
+        df = pandas.read_csv(dataset_params["dataset_description"])
+    training_df, validation_df = train_test_split(df, test_size=dataset_params["validation_ratio"])
+    torch.manual_seed(dataset_params['seed'])
+    training_set = SideSet(dataset_yaml,
+                           set_type="train",
+                           dataset_df=training_df,
+                           chunk_per_segment=dataset_params['train']['chunk_per_segment'],
+                           overlap=dataset_params['train']['overlap'])
+
+    validation_set = SideSet(dataset_yaml, set_type="validation", dataset_df=validation_df)
+
+    return model, training_set, validation_set
+
+def xtime(model, training_set, validation_set,
+          batch_size,
+          num_thread,
+          ):
+    training_loader = DataLoader(training_set,
+                                 batch_size=batch_size,
+                                 shuffle=True,
+                                 drop_last=True,
+                                 pin_memory=True,
+                                 num_workers=num_thread)
+
+    validation_loader = DataLoader(validation_set,
+                                   batch_size=batch_size,
+                                   drop_last=True,
+                                   pin_memory=True,
+                                   num_workers=num_thread)
+    """
+    Set the training options
+    """
+    if opt == 'adam':
+        _optimizer = torch.optim.Adam
+        _options = {'lr': lr}
+    elif opt == 'rmsprop':
+        _optimizer = torch.optim.RMSprop
+        _options = {'lr': lr}
+    else:  # opt == 'sgd'
+        _optimizer = torch.optim.SGD
+        _options = {'lr': lr, 'momentum': 0.9}
+
+    params = [
+        {
+            'params': [
+                param for name, param in model.named_parameters() if 'bn' not in name
+            ]
+        },
+        {
+            'params': [
+                param for name, param in model.named_parameters() if 'bn' in name
+            ],
+            'weight_decay': 0
+        },
+    ]
+
+    param_list = []
+    if type(model) is Xtractor:
+        if model.preprocessor is not None:
+            param_list.append(
+                {'params': model.preprocessor.parameters(), 'weight_decay': model.preprocessor_weight_decay})
+        param_list.append(
+            {'params': model.sequence_network.parameters(), 'weight_decay': model.sequence_network_weight_decay})
+        param_list.append({'params': model.stat_pooling.parameters(), 'weight_decay': model.stat_pooling_weight_decay})
+        param_list.append({'params': model.before_speaker_embedding.parameters(),
+                           'weight_decay': model.before_speaker_embedding_weight_decay})
+        param_list.append({'params': model.after_speaker_embedding.parameters(),
+                           'weight_decay': model.after_speaker_embedding_weight_decay})
+
+    else:
+        if model.module.preprocessor is not None:
+            param_list.append({'params': model.module.preprocessor.parameters(),
+                               'weight_decay': model.module.preprocessor_weight_decay})
+        param_list.append({'params': model.module.sequence_network.parameters(),
+                           'weight_decay': model.module.sequence_network_weight_decay})
+        param_list.append(
+            {'params': model.module.stat_pooling.parameters(), 'weight_decay': model.module.stat_pooling_weight_decay})
+        param_list.append({'params': model.module.before_speaker_embedding.parameters(),
+                           'weight_decay': model.module.before_speaker_embedding_weight_decay})
+        param_list.append({'params': model.module.after_speaker_embedding.parameters(),
+                           'weight_decay': model.module.after_speaker_embedding_weight_decay})
+
+    optimizer = _optimizer(param_list, **_options)
+
+    # optimizer = torch.optim.SGD(params,
+    #                            lr=lr,
+    #                            momentum=0.9,
+    #                            weight_decay=0.0005)
+    # print(f"Learning rate = {lr}")
+
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', verbose=True)
+
+    best_accuracy = 0.0
+    best_accuracy_epoch = 1
+    curr_patience = patience
+    for epoch in range(1, epochs + 1):
+        # Process one epoch and return the current model
+        if curr_patience == 0:
+            print(f"Stopping at epoch {epoch} for cause of patience")
+            break
+        model = train_epoch(model,
+                            epoch,
+                            training_loader,
+                            optimizer,
+                            dataset_params["log_interval"],
+                            device=device,
+                            clipping=clipping,
+                            tb_writer=writer)
+
+        # Add the cross validation here
+        accuracy, val_loss = cross_validation(model, validation_loader, device=device)
+        logging.critical(f"***{time.strftime('%H:%M:%S', time.localtime())} Cross validation accuracy = {accuracy} %")
+
+        # Decrease learning rate according to the scheduler policy
+        scheduler.step(val_loss)
+
+        # remember best accuracy and save checkpoint
+        is_best = accuracy > best_accuracy
+        best_accuracy = max(accuracy, best_accuracy)
+
+        if type(model) is Xtractor:
+            save_checkpoint({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'accuracy': best_accuracy,
+                'scheduler': scheduler,
+                'speaker_number': speaker_number,
+                'model_archi': model_archi
+            }, is_best, filename=tmp_model_name + ".pt", best_filename=best_model_name + '.pt')
+        else:
+            save_checkpoint({
+                'epoch': epoch,
+                'model_state_dict': model.module.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'accuracy': best_accuracy,
+                'scheduler': scheduler,
+                'speaker_number': speaker_number,
+                'model_archi': model_archi
+            }, is_best, filename=tmp_model_name + ".pt", best_filename=best_model_name + '.pt')
+
+        if is_best:
+            best_accuracy_epoch = epoch
+            curr_patience = patience
+        else:
+            curr_patience -= 1
+    # writer.close()
+
+    for ii in range(torch.cuda.device_count()):
+        print(torch.cuda.memory_summary(ii))
+
+    logging.critical(f"Best accuracy {best_accuracy * 100.} obtained at epoch {best_accuracy_epoch}")