xvector.py

# coding: utf-8 -*-
#
# This file is part of SIDEKIT.
#
# SIDEKIT is a python package for speaker verification.
# Home page: http://www-lium.univ-lemans.fr/sidekit/
#
# SIDEKIT is a python package for speaker verification.
# Home page: http://www-lium.univ-lemans.fr/sidekit/
#
# SIDEKIT is free software: you can redistribute it and/or modify
# it under the terms of the GNU LLesser General Public License as
# published by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# SIDEKIT is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with SIDEKIT.  If not, see <http://www.gnu.org/licenses/>.

"""
Copyright 2014-2021 Anthony Larcher, Yevhenii Prokopalo
"""


import pdb
import traceback
import logging
import math
import os
import numpy
import pandas
import pickle
import shutil
import sys
import time
import torch
import torchaudio
import tqdm
import yaml

from torchvision import transforms
from collections import OrderedDict
from .xsets import SideSet
from .xsets import FileSet
from .xsets import IdMapSet
from .xsets import IdMapSet_per_speaker
from .xsets import SpkSet
from .res_net import RawPreprocessor, ResBlockWFMS, ResBlock, PreResNet34, PreFastResNet34
from ..bosaris import IdMap
from ..bosaris import Key
from ..bosaris import Ndx
from ..bosaris.detplot import rocch
from ..bosaris.detplot import rocch2eer
from ..statserver import StatServer
from ..iv_scoring import cosine_scoring
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split
from .sincnet import SincNet
from .loss import ArcLinear
from .loss import ArcFace
from .loss import l2_norm
from .loss import ArcMarginProduct


os.environ['MKL_THREADING_LAYER'] = 'GNU'

__license__ = "LGPL"
__author__ = "Anthony Larcher"
__copyright__ = "Copyright 2015-2021 Anthony Larcher"
__maintainer__ = "Anthony Larcher"
__email__ = "anthony.larcher@univ-lemans.fr"
__status__ = "Production"
__docformat__ = 'reS'


logging.basicConfig(format='%(asctime)s %(message)s')


# Make PyTorch Deterministic
torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
numpy.random.seed(0)


class GuruMeditation (torch.autograd.detect_anomaly):
    
    def __init__(self):
        super(GuruMeditation, self).__init__()

    def __enter__(self):
        super(GuruMeditation, self).__enter__()
        return self

    def __exit__(self, type, value, trace):
        super(GuruMeditation, self).__exit__()
        if isinstance(value, RuntimeError):
            traceback.print_tb(trace)
            self.halt(str(value))

    def halt(msg):
        print (msg)
        pdb.set_trace()


def select_n_random(data, labels, n=100):
    '''
    Selects n random datapoints and their corresponding labels from a dataset
    '''
    assert len(data) == len(labels)

    perm = torch.randperm(len(data))
    return data[perm][:n], labels[perm][:n]


def matplotlib_imshow(img, one_channel=False):
    if one_channel:
        img = img.mean(dim=0)
    img = img / 2 + 0.5     # unnormalize
    npimg = img.cpu().numpy()
    if one_channel:
        plt.imshow(npimg, cmap="Greys")
    else:
        plt.imshow(numpy.transpose(npimg, (1, 2, 0)))

def speech_to_probs(model, speech):
    '''
    Generates predictions and corresponding probabilities from a trained
    network and a list of images
    '''
    output = model(speech)
    # convert output probabilities to predicted class
    _, preds_tensor = torch.max(output, 1)
    preds = numpy.squeeze(preds_tensor.cpu().numpy())
    return preds, [torch.nn.functional.softmax(el, dim=0)[i].item() for i, el in zip(preds, output)]


def plot_classes_preds(model, speech, labels):
    '''

    Generates matplotlib Figure using a trained network, along with images
    and labels from a batch, that shows the network's top prediction along
    with its probability, alongside the actual label, coloring this
    information based on whether the prediction was correct or not.
    Uses the "speech_to_probs" function.
    '''
    preds, probs = speech_to_probs(model, speech)
    # plot the images in the batch, along with predicted and true labels
    fig = plt.figure(figsize=(12, 48))
    for idx in numpy.arange(4):
         ax = fig.add_subplot(1, 4, idx+1, xticks=[], yticks=[])
         #matplotlib_imshow(speech[idx], one_channel=True)
         ax.set_title("{0}, {1:.1f}%\n(label: {2})".format(
                                                           preds[idx],
                                                           probs[idx] * 100.0,
                                                           labels[idx]),
                                                           color=("green" if preds[idx]==labels[idx].item() else "red"))
    return fig


def test_metrics(model,
                 device,
                 speaker_number,
                 num_thread,
                 mixed_precision):
    """Compute model metrics

    Args:
        model ([type]): [description]
        validation_loader ([type]): [description]
        device ([type]): [description]
        speaker_number ([type]): [description]
        model_archi ([type]): [description]

    Raises:
        NotImplementedError: [description]
        NotImplementedError: [description]

    Returns:
        [type]: [description]
    """
    idmap_test_filename = 'h5f/idmap_test.h5'
    ndx_test_filename = 'h5f/ndx_test.h5'
    key_test_filename = 'h5f/key_test.h5'
    data_root_name='/lium/corpus/base/voxceleb1/test/wav'

    transform_pipeline = dict()
    #mfcc_config = dict()
    #mfcc_config['nb_filters'] = 81
    #mfcc_config['nb_ceps'] = 80
    #mfcc_config['lowfreq'] = 133.333
    #mfcc_config['maxfreq'] = 6855.4976
    #mfcc_config['win_time'] = 0.025
    #mfcc_config['shift'] = 0.01
    #mfcc_config['n_fft'] = 2048
    #transform_pipeline['MFCC'] = mfcc_config
    #transform_pipeline['CMVN'] = {}

    xv_stat = extract_embeddings(idmap_name=idmap_test_filename,
                                 speaker_number=speaker_number,
                                 model_filename=model,
                                 data_root_name=data_root_name,
                                 device=device,
                                 transform_pipeline=transform_pipeline,
                                 num_thread=num_thread,
                                 mixed_precision=mixed_precision)

    scores = cosine_scoring(xv_stat,
                            xv_stat,
                            Ndx(ndx_test_filename),
                            wccn=None,
                            check_missing=True)

    tar, non = scores.get_tar_non(Key(key_test_filename))

    test_eer = eer(numpy.array(non).astype(numpy.double), numpy.array(tar).astype(numpy.double))

    return test_eer


def get_lr(optimizer):
    """

    :param optimizer:
    :return:
    """
    for param_group in optimizer.param_groups:
        return param_group['lr']


def save_checkpoint(state, is_best, filename='checkpoint.pth.tar', best_filename='model_best.pth.tar'):
    """

    :param state:
    :param is_best:
    :param filename:
    :param best_filename:
    :return:
    """
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, best_filename)

class MeanStdPooling(torch.nn.Module):
    """
    Mean and Standard deviation pooling
    """
    def __init__(self):
        """

        """
        super(MeanStdPooling, self).__init__()
        pass

    def forward(self, x):
        """

        :param x:
        :return:
        """
        mean = torch.mean(x, dim=2)
        std = torch.std(x, dim=2)
        return torch.cat([mean, std], dim=1)


class GruPooling(torch.nn.Module):
    """

    """
    def __init__(self, input_size, gru_node, nb_gru_layer):
        """

        :param input_size:
        :param gru_node:
        :param nb_gru_layer:
        """
        super(GruPooling, self).__init__()
        self.lrelu_keras = torch.nn.LeakyReLU(negative_slope = 0.3)
        self.bn_before_gru = torch.nn.BatchNorm1d(num_features = input_size)
        self.gru = torch.nn.GRU(input_size = input_size,
                          hidden_size = gru_node,
                          num_layers = nb_gru_layer,
                          batch_first = True)

    def forward(self, x):
        """

        :param x:
        :return:
        """
        x = self.bn_before_gru(x)
        x = self.lrelu_keras(x)
        x = x.permute(0, 2, 1)  #(batch, filt, time) >> (batch, time, filt)
        self.gru.flatten_parameters()
        x, _ = self.gru(x)
        x = x[:,-1,:]

        return x


class PreEmphasis(torch.nn.Module):

    def __init__(self, coef: float = 0.97):
        super().__init__()
        self.coef = coef
        # make kernel
        # In pytorch, the convolution operation uses cross-correlation. So, filter is flipped.
        self.register_buffer(
            'flipped_filter', torch.FloatTensor([-self.coef, 1.]).unsqueeze(0).unsqueeze(0)
        )

    def forward(self, input: torch.tensor) -> torch.tensor:
        assert len(input.size()) == 2, 'The number of dimensions of input tensor must be 2!'
        # reflect padding to match lengths of in/out
        input = input.unsqueeze(1)
        input = torch.nn.functional.pad(input, (1, 0), 'reflect')
        return torch.nn.functional.conv1d(input, self.flipped_filter).squeeze(1)


class MfccFrontEnd(torch.nn.Module):
    """

    """

    def __init__(self,
                 pre_emphasis=0.97,
                 sample_rate=16000,
                 n_fft=2048,
                 f_min=133.333,
                 f_max=6855.4976,
                 win_length=1024,
                 window_fn=torch.hann_window,
                 hop_length=512,
                 power=2.0,
                 n_mels=100,
                 n_mfcc=80):

        super(MfccFrontEnd, self).__init__()

        self.pre_emphasis = pre_emphasis
        self.sample_rate = sample_rate
        self.n_fft = n_fft
        self.f_min = f_min
        self.f_max = f_max
        self.win_length = win_length
        self.window_fn=window_fn
        self.hop_length = hop_length
        self.power=power
        self.window_fn = window_fn
        self.n_mels = n_mels
        self.n_mfcc = n_mfcc

        self.PreEmphasis = PreEmphasis(self.pre_emphasis)

        self.melkwargs = {"n_fft":self.n_fft,
                          "f_min":self.f_min,
                          "f_max":self.f_max,
                          "win_length":self.win_length,
                          "window_fn":self.window_fn,
                          "hop_length":self.hop_length,
                          "power":self.power,
                          "n_mels":self.n_mels}

        self.MFCC = torchaudio.transforms.MFCC(
            sample_rate=self.sample_rate,
            n_mfcc=self.n_mfcc,
            dct_type=2,
            log_mels=True,
            melkwargs=self.melkwargs)

        self.CMVN = torch.nn.InstanceNorm1d(self.n_mfcc)

    def forward(self, x):
        """

        :param x:
        :return:
        """
        with torch.no_grad():
            with torch.cuda.amp.autocast(enabled=False):
                mfcc = self.PreEmphasis(x)
                mfcc = self.MFCC(mfcc)
                mfcc = self.CMVN(mfcc)
        return mfcc

class Xtractor(torch.nn.Module):
    """
    Class that defines an x-vector extractor based on 5 convolutional layers and a mean standard deviation pooling
    """

    def __init__(self,
                 speaker_number,
                 model_archi="xvector",
                 loss=None,
                 norm_embedding=False,
                 aam_margin=0.2,
                 aam_s=30):
        """
        If config is None, default architecture is created
        :param model_archi:
        """
        super(Xtractor, self).__init__()
        self.speaker_number = speaker_number
        self.feature_size = None
        self.norm_embedding = norm_embedding

        print(f"Speaker number : {self.speaker_number}")

        if model_archi == "xvector":

            self.input_nbdim = 2

            if loss not in ["cce", 'aam']:
                raise NotImplementedError(f"The valid loss are for now cce and aam ")
            else:
                self.loss = loss

            self.activation = torch.nn.LeakyReLU(0.2)

            self.preprocessor = MfccFrontEnd()
            self.feature_size = self.preprocessor.n_mfcc

            self.sequence_network = torch.nn.Sequential(OrderedDict([
                ("conv1", torch.nn.Conv1d(self.feature_size, 512, 5, dilation=1)),
                ("activation1", torch.nn.LeakyReLU(0.2)),
                ("batch_norm1", torch.nn.BatchNorm1d(512)),
                ("conv2", torch.nn.Conv1d(512, 512, 3, dilation=2)),
                ("activation2", torch.nn.LeakyReLU(0.2)),
                ("batch_norm2", torch.nn.BatchNorm1d(512)),
                ("conv3", torch.nn.Conv1d(512, 512, 3, dilation=3)),
                ("activation3", torch.nn.LeakyReLU(0.2)),
                ("batch_norm3", torch.nn.BatchNorm1d(512)),
                ("conv4", torch.nn.Conv1d(512, 512, 1)),
                ("activation4", torch.nn.LeakyReLU(0.2)),
                ("batch_norm4", torch.nn.BatchNorm1d(512)),
                ("conv5", torch.nn.Conv1d(512, 1536, 1)),
                ("activation5", torch.nn.LeakyReLU(0.2)),
                ("batch_norm5", torch.nn.BatchNorm1d(1536))
            ]))

            self.stat_pooling = MeanStdPooling()
            self.stat_pooling_weight_decay = 0
            self.before_speaker_embedding = torch.nn.Sequential(OrderedDict([
                ("linear6", torch.nn.Linear(3072, 512))
            ]))

            self.embedding_size = 512

            if self.loss == "aam":
                self.after_speaker_embedding = ArcMarginProduct(512,
                                                                int(self.speaker_number),
                                                                s=64,
                                                                m=0.2,
                                                                easy_margin=True)
            elif self.loss == "cce":
                self.after_speaker_embedding = torch.nn.Sequential(OrderedDict([
                    ("activation6", torch.nn.LeakyReLU(0.2)),
                    ("batch_norm6", torch.nn.BatchNorm1d(512)),
                    ("dropout6", torch.nn.Dropout(p=0.05)),
                    ("linear7", torch.nn.Linear(512, 512)),
                    ("activation7", torch.nn.LeakyReLU(0.2)),
                    ("batch_norm7", torch.nn.BatchNorm1d(512)),
                    ("linear8", torch.nn.Linear(512, int(self.speaker_number)))
                ]))

            self.sequence_network_weight_decay = 0.0002
            self.before_speaker_embedding_weight_decay = 0.002
            self.after_speaker_embedding_weight_decay = 0.002
            self.embedding_size = 512

        elif model_archi == "resnet34":

            self.preprocessor = MfccFrontEnd()
            self.sequence_network = PreResNet34()

            self.before_speaker_embedding = torch.nn.Linear(in_features = 5120,
                                                            out_features = 256)

            self.stat_pooling = MeanStdPooling()
            self.stat_pooling_weight_decay = 0

            self.embedding_size = 256

            self.loss = "aam"
            self.after_speaker_embedding = ArcMarginProduct(256,
                                                            int(self.speaker_number),
                                                            s = 30.0,
                                                            m = 0.20,
                                                            easy_margin = True)

            self.preprocessor_weight_decay = 0.000
            self.sequence_network_weight_decay = 0.000
            self.stat_pooling_weight_decay = 0.000
            self.before_speaker_embedding_weight_decay = 0.00
            self.after_speaker_embedding_weight_decay = 0.00

        elif model_archi == "fastresnet34":

            self.preprocessor = MfccFrontEnd()
            self.sequence_network = PreFastResNet34()
            self.embedding_size = 256

            self.before_speaker_embedding = torch.nn.Linear(in_features = 2560,
                                                            out_features = 256)

            self.stat_pooling = MeanStdPooling()
            self.stat_pooling_weight_decay = 0

            self.loss = "aam"
            self.after_speaker_embedding = ArcMarginProduct(self.embedding_size,
                                                            int(self.speaker_number),
                                                            s = 30,
                                                            m = 0.2,
                                                            easy_margin = False)

            self.preprocessor_weight_decay = 0.000
            self.sequence_network_weight_decay = 0.000
            self.stat_pooling_weight_decay = 0.000
            self.before_speaker_embedding_weight_decay = 0.00
            self.after_speaker_embedding_weight_decay = 0.00

        elif model_archi == "rawnet2":

            if loss not in ["cce", 'aam']:
                raise NotImplementedError(f"The valid loss are for now cce and aam ")
            else:
                self.loss = loss

            self.input_nbdim = 2

            filts = [128, [128, 128], [128, 256], [256, 256]]
            self.norm_embedding = True

            self.preprocessor = RawPreprocessor(nb_samp=48000,
                                                in_channels=1,
                                                out_channels=filts[0],
                                                kernel_size=3)

            self.sequence_network = torch.nn.Sequential(OrderedDict([
                ("block0", ResBlockWFMS(nb_filts=filts[1], first=True)),
                ("block1", ResBlockWFMS(nb_filts=filts[1])),
                ("block2", ResBlockWFMS(nb_filts=filts[2])),
                ("block3", ResBlockWFMS(nb_filts=[filts[2][1], filts[2][1]])),
                ("block4", ResBlockWFMS(nb_filts=[filts[2][1], filts[2][1]])),
                ("block5", ResBlockWFMS(nb_filts=[filts[2][1], filts[2][1]]))
            ]))

            self.stat_pooling = GruPooling(input_size=filts[2][-1],
                                           gru_node=1024,
                                           nb_gru_layer=1)

            self.before_speaker_embedding = torch.nn.Linear(in_features = 1024,
                                                            out_features = 1024)

            if self.loss == "aam":
                if loss == 'aam':
                    self.after_speaker_embedding = ArcLinear(1024,
                                                             int(self.speaker_number),
                                                             margin=aam_margin, s=aam_s)
            elif self.loss == "cce":
                self.after_speaker_embedding = torch.nn.Linear(in_features = 1024,
                                                               out_features = int(self.speaker_number),
                                                               bias = True)

            self.preprocessor_weight_decay = 0.000
            self.sequence_network_weight_decay = 0.000
            self.stat_pooling_weight_decay = 0.000
            self.before_speaker_embedding_weight_decay = 0.00
            self.after_speaker_embedding_weight_decay = 0.00

        else:
            is_first_resblock = True

            if isinstance(model_archi, dict):
                cfg = model_archi
            else:
                # Load Yaml configuration
                with open(model_archi, 'r') as fh:
                    cfg = yaml.load(fh, Loader=yaml.FullLoader)

            self.loss = cfg["training"]["loss"]
            if self.loss == "aam":
                self.aam_margin = cfg["training"]["aam_margin"]
                self.aam_s = cfg["training"]["aam_s"]

            """
            Prepare Preprocessor
            """
            self.preprocessor = None
            if "preprocessor" in cfg:
                if cfg['preprocessor']["type"] == "sincnet":
                    self.preprocessor = SincNet(
                        waveform_normalize=cfg['preprocessor']["waveform_normalize"],
                        sample_rate=cfg['preprocessor']["sample_rate"],
                        min_low_hz=cfg['preprocessor']["min_low_hz"],
                        min_band_hz=cfg['preprocessor']["min_band_hz"],
                        out_channels=cfg['preprocessor']["out_channels"],
                        kernel_size=cfg['preprocessor']["kernel_size"],
                        stride=cfg['preprocessor']["stride"],
                        max_pool=cfg['preprocessor']["max_pool"],
                        instance_normalize=cfg['preprocessor']["instance_normalize"],
                        activation=cfg['preprocessor']["activation"],
                        dropout=cfg['preprocessor']["dropout"]
                    )
                    self.feature_size = self.preprocessor.dimension
                elif cfg['preprocessor']["type"] == "rawnet2":
                    self.preprocessor = RawPreprocessor(nb_samp=int(cfg['preprocessor']["sampling_frequency"] * cfg['preprocessor']["duration"]),
                                                        in_channels=1,
                                                        out_channels=cfg["feature_size"],
                                                        kernel_size=cfg['preprocessor']["kernel_size"],
                                                        stride=cfg['preprocessor']["stride"],
                                                        padding=cfg['preprocessor']["padding"],
                                                        dilation=cfg['preprocessor']["dilation"])
                    self.feature_size = cfg["feature_size"]
                self.preprocessor_weight_decay = 0.000

            """
            Prepare sequence network
            """
            # Get Feature size
            if self.feature_size is None:
                self.feature_size = cfg["feature_size"]

            input_size = self.feature_size

            # Get activation function
            if cfg["activation"] == 'LeakyReLU':
                self.activation = torch.nn.LeakyReLU(0.2)
            elif cfg["activation"] == 'PReLU':
                self.activation = torch.nn.PReLU()
            elif cfg["activation"] == 'ReLU6':
                self.activation = torch.nn.ReLU6()
            else:
                self.activation = torch.nn.ReLU()

            # Create sequential object for the first part of the network
            segmental_layers = []
            for k in cfg["segmental"].keys():
                if k.startswith("lin"):
                    segmental_layers.append((k, torch.nn.Linear(input_size,
                                                                cfg["segmental"][k]["output"])))
                    input_size = cfg["segmental"][k]["output"]

                elif k.startswith("conv2D"):
                    segmental_layers.append((k, torch.nn.Conv2d(in_channels=1,
                                                                out_channels=entry_conv_out_channels,
                                                                kernel_size=entry_conv_kernel_size,
                                                                padding=3,
                                                                stride=1)))

                elif k.startswith("conv"):
                    segmental_layers.append((k, torch.nn.Conv1d(input_size,
                                                                cfg["segmental"][k]["output_channels"],
                                                                kernel_size=cfg["segmental"][k]["kernel_size"],
                                                                dilation=cfg["segmental"][k]["dilation"])))
                    input_size = cfg["segmental"][k]["output_channels"]

                elif k.startswith("ctrans"):
                    segmental_layers.append((k, torch.nn.ConvTranspose1d(input_size,
                                                                         cfg["segmental"][k]["output_channels"],
                                                                         kernel_size=cfg["segmental"][k]["kernel_size"],
                                                                         dilation=cfg["segmental"][k]["dilation"])))
                elif k.startswith("activation"):
                    segmental_layers.append((k, self.activation))

                elif k.startswith('batch_norm'):
                    segmental_layers.append((k, torch.nn.BatchNorm1d(input_size)))

                elif k.startswith('resblock'):
                    segmental_layers.append((ResBlock(cfg["segmental"][k]["input_channel"],
                                                      cfg["segmental"][k]["output_channel"],
                                                      is_first_resblock)))
                    is_first_resblock = False

            self.sequence_network = torch.nn.Sequential(OrderedDict(segmental_layers))
            self.sequence_network_weight_decay = cfg["segmental"]["weight_decay"]

            """
            Pooling
            """
            self.stat_pooling = MeanStdPooling()
            tmp_input_size = input_size * 2
            if cfg["stat_pooling"]["type"] == "GRU":
                self.stat_pooling = GruPooling(input_size=cfg["stat_pooling"]["input_size"],
                                               gru_node=cfg["stat_pooling"]["gru_node"],
                                               nb_gru_layer=cfg["stat_pooling"]["nb_gru_layer"])
                tmp_input_size = cfg["stat_pooling"]["gru_node"]

            self.stat_pooling_weight_decay = cfg["stat_pooling"]["weight_decay"]

            """
            Prepare last part of the network (after pooling)
            """
            # Create sequential object for the second part of the network
            input_size = tmp_input_size
            before_embedding_layers = []
            for k in cfg["before_embedding"].keys():
                if k.startswith("lin"):
                    if cfg["before_embedding"][k]["output"] == "speaker_number":
                        before_embedding_layers.append((k, torch.nn.Linear(input_size, self.speaker_number)))
                    else:
                        before_embedding_layers.append((k, torch.nn.Linear(input_size,
                                                                           cfg["before_embedding"][k]["output"])))
                        input_size = cfg["before_embedding"][k]["output"]

                elif k.startswith("activation"):
                    before_embedding_layers.append((k, self.activation))

                elif k.startswith('batch_norm'):
                    before_embedding_layers.append((k, torch.nn.BatchNorm1d(input_size)))

                elif k.startswith('dropout'):
                    before_embedding_layers.append((k, torch.nn.Dropout(p=cfg["before_embedding"][k])))

            self.embedding_size = input_size
            self.before_speaker_embedding = torch.nn.Sequential(OrderedDict(before_embedding_layers))
            self.before_speaker_embedding_weight_decay = cfg["before_embedding"]["weight_decay"]

            # if loss_criteria is "cce"
            # Create sequential object for the second part of the network
            if self.loss == "cce":
                after_embedding_layers = []
                for k in cfg["after_embedding"].keys():
                    if k.startswith("lin"):
                        if cfg["after_embedding"][k]["output"] == "speaker_number":
                            after_embedding_layers.append((k, torch.nn.Linear(input_size, self.speaker_number)))
                        else:
                            after_embedding_layers.append((k, torch.nn.Linear(input_size,
                                                                          cfg["after_embedding"][k]["output"])))
                            input_size = cfg["after_embedding"][k]["output"]

                    elif k.startswith('arc'):
                        after_embedding_layers.append((k, ArcLinear(input_size,
                                                                    self.speaker_number,
                                                                    margin=self.aam_margin,
                                                                    s=self.aam_s)))

                    elif k.startswith("activation"):
                        after_embedding_layers.append((k, self.activation))

                    elif k.startswith('batch_norm'):
                        after_embedding_layers.append((k, torch.nn.BatchNorm1d(input_size)))

                    elif k.startswith('dropout'):
                        after_embedding_layers.append((k, torch.nn.Dropout(p=cfg["after_embedding"][k])))

                self.after_speaker_embedding = torch.nn.Sequential(OrderedDict(after_embedding_layers))

            elif self.loss == "aam":
                self.norm_embedding = True
                self.after_speaker_embedding = ArcMarginProduct(input_size,
                                                                int(self.speaker_number),
                                                                s=64,
                                                                m=0.2,
                                                                easy_margin=True)

                #self.after_speaker_embedding = ArcLinear(input_size,
                #                                         self.speaker_number,
                #                                         margin=self.aam_margin,
                #                                         s=self.aam_s)
                #self.after_speaker_embedding = ArcFace(embedding_size=input_size,
                #                                       classnum=self.speaker_number,
                #                                       s=64.,
                #                                       m=0.5)

            self.after_speaker_embedding_weight_decay = cfg["after_embedding"]["weight_decay"]


    def forward(self, x, is_eval=False, target=None, extract_after_pooling=False):
        """

        :param x:
        :param is_eval: False for training
        :return:
        """
        if self.preprocessor is not None:
            x = self.preprocessor(x)

        x = self.sequence_network(x)

        # Mean and Standard deviation pooling
        x = self.stat_pooling(x)

        if extract_after_pooling:
            return x

        x = self.before_speaker_embedding(x)

        if self.norm_embedding:
            #x_norm = x.norm(p=2,dim=1, keepdim=True) / 10. # Why  10. ?
            #x_norm = torch.linalg.norm(x, ord=2, dim=1, keepdim=True, out=None, dtype=None)
            #x = torch.div(x, x_norm)
            x = l2_norm(x)

        if self.loss == "cce":
            if is_eval:
                return x
            else:
                return self.after_speaker_embedding(x), x

        elif self.loss == "aam":
            if is_eval:
                x = torch.nn.functional.normalize(x, dim=1)
            else:
                x = self.after_speaker_embedding(torch.nn.functional.normalize(x, dim=1), target=target), torch.nn.functional.normalize(x, dim=1)

        return x

    def context_size(self):
        context = 1
        if isinstance(self, Xtractor):
            for name, module in self.sequence_network.named_modules():
                if name.startswith("conv"):
                    context += module.dilation[0] * (module.kernel_size[0] - 1)
        else:
            for name, module in self.module.sequence_network.named_modules():
                if name.startswith("conv"):
                    context += module.dilation[0] * (module.kernel_size[0] - 1)
        return context

def xtrain(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,
           mixed_precision=False,
           clipping=False,
           opt=None,
           reset_parts=[],
           freeze_parts=[],
           num_thread=None,
           write_batches_to_disk=False,
           load_batches_from_disk=False,
           tmp_batch_dir=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 write_batches_to_disk:
        load_batches_from_disk = True

    if num_thread is None:
        import multiprocessing

        num_thread = multiprocessing.cpu_count()

    logging.critical(f"Use {num_thread} cpus")
    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", "resnet34", "fastresnet34"]:
        # Initialize a first model
        if model_yaml == "xvector":
            model = Xtractor(speaker_number, "xvector", loss=loss)
        elif model_yaml == "rawnet2":
            model = Xtractor(speaker_number, "rawnet2")
        elif model_yaml == "resnet34":
            model = Xtractor(speaker_number, "resnet34")
        elif model_yaml == "fastresnet34":
            model = Xtractor(speaker_number, "fastresnet34")
        model_archi = model_yaml
    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

    logging.critical(model)

    logging.critical("model_parameters_count: {:d}".format(
        sum(p.numel()
            for p in model.sequence_network.parameters()
            if p.requires_grad) + \
        sum(p.numel()
            for p in model.before_speaker_embedding.parameters()
            if p.requires_grad)))

    embedding_size = model.embedding_size

    if torch.cuda.device_count() > 1 and multi_gpu:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = torch.nn.DataParallel(model)
        #model = DDP(model)
    else:
        print("Train on a single GPU")
    model.to(device)

    with open(dataset_yaml, "r") as fh:
        dataset_params = yaml.load(fh, Loader=yaml.FullLoader)
        df = pandas.read_csv(dataset_params["dataset_description"])

    if load_batches_from_disk:
        train_batch_fn_format = tmp_batch_dir + "/train/train_{}_batch.h5"
        val_batch_fn_format = tmp_batch_dir + "/val/val_{}_batch.h5"

    if not load_batches_from_disk or write_batches_to_disk:
        """
        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 
        """

        if write_batches_to_disk or dataset_params["batch_size"] > 1:
            output_format = "numpy"
        else:
            output_format = "pytorch"

        training_df, validation_df = train_test_split(df, test_size=dataset_params["validation_ratio"])
        torch.manual_seed(dataset_params['seed'])
        training_set = SpkSet(dataset_yaml,
                              set_type="train",
                              dataset_df=training_df,
                              overlap=dataset_params['train']['overlap'],
                              output_format=output_format,
                              windowed=True)