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 logging
import math
import os
import numpy
import pandas
import pickle
import shutil
import time
import torch
import torchaudio
import tqdm
import yaml

from collections import OrderedDict
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split
from .preprocessor import MfccFrontEnd
from .preprocessor import MelSpecFrontEnd
from .preprocessor import RawPreprocessor
from .xsets import SideSet
from .xsets import IdMapSet
from .xsets import IdMapSetPerSpeaker
from .xsets import SideSampler
from .res_net import ResBlockWFMS
from .res_net import ResBlock
from .res_net import PreResNet34
from .res_net import PreFastResNet34
from ..bosaris import IdMap
from ..bosaris import Key
from ..bosaris import Ndx
from ..statserver import StatServer
from ..iv_scoring import cosine_scoring
from .sincnet import SincNet
from .loss import ArcLinear
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)



def eer(negatives, positives):
    """Logarithmic complexity EER computation

    Args:
        negative_scores (numpy array): impostor scores
        positive_scores (numpy array): genuine scores

    Returns:
        float: Equal Error Rate (EER)
    """

    positives = numpy.sort(positives)
    negatives = numpy.sort(negatives)[::-1]

    pos_count = positives.shape[0]
    neg_count = negatives.shape[0]

    p_score = positives[0]
    n_score = negatives[0]

    p_index = 0
    n_index = 0

    next_p_jump = pos_count//2
    next_n_jump = neg_count//2

    kdx = 0
    while True:
        kdx += 1
        if p_index < 0 or n_index < 0:
            return 0
        if p_index > pos_count or n_index > neg_count:
            return 100
        if p_score < n_score:
            p_index = p_index + next_p_jump
            n_index = n_index + next_n_jump
            if next_p_jump == 0 and next_n_jump == 0:
                break
        elif p_score >= n_score:
            p_index = p_index - next_p_jump
            n_index = n_index - next_n_jump
            if next_p_jump == 0 and next_n_jump == 0:
                break

        p_score = positives[p_index]
        n_score = negatives[n_index]
        next_p_jump = next_p_jump//2
        next_n_jump = next_n_jump//2

    eer_predicate = 100

    tfr = (abs(p_index))/pos_count
    tfa = (1+abs(n_index))/neg_count
    if (p_score == n_score and tfr == tfa):
        return tfr

    while positives[p_index] < negatives[n_index]:
        if p_index < pos_count - 1:
            p_index += 1
        elif n_index < neg_count - 1:
            n_index += 1
        else:
            break

    while positives[p_index] > negatives[n_index] and n_index >= 1:
        n_index -= 1

    tfr = (1+p_index)/pos_count
    tfa = (1+n_index)/neg_count

    while tfa > tfr:
        p_index += 1
        while positives[p_index] > negatives[n_index] and n_index >= 1:
            n_index -= 1
        tfr = (1+p_index)/pos_count
        tfa = (1+n_index)/neg_count

    if abs(tfr - tfa) <= eer_predicate:
        eer_predicate = abs(tfr - tfa)
        eer = (tfr + tfa) / 2
    else:
        return eer

    tfr = p_index/pos_count
    tfa = (1+n_index)/neg_count
    if abs(tfr - tfa) <= eer_predicate:
        eer_predicate = abs(tfr - tfa)
        eer = (tfr + tfa) / 2
    else:
        return eer

    while True:
        while negatives[n_index + 1] <= positives[p_index - 1]:
            p_index -= 1
            tfr = p_index/pos_count
            tfa = (1+n_index)/neg_count
            if abs(tfr - tfa) <= eer_predicate:
                eer_predicate = abs(tfr - tfa)
                eer = (tfr + tfa) / 2
            else:
                return eer
        while negatives[n_index + 1] > positives[p_index - 1]:
            n_index += 1
            tfr = p_index/pos_count
            tfa = (1+n_index)/neg_count
            if abs(tfr - tfa) <= eer_predicate:
                eer_predicate = abs(tfr - tfa)
                eer = (tfr + tfa) / 2
            else:
                return eer

    return eer


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='/data/larcher/voxceleb1/test/wav'

    transform_pipeline = dict()

    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 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 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 = MelSpecFrontEnd()
            self.sequence_network = PreFastResNet34()
            self.embedding_size = 256

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

            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,
           compute_test_eer=True):
    """

    :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 mixed_precision:
    :param clipping:
    :param opt:
    :param reset_parts:
    :param freeze_parts:
    :param num_thread:
    :param compute_test_eer:
    :return:
    """
    # Test to optimize
    torch.autograd.profiler.emit_nvtx(enabled=False)

    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" if torch.cuda.is_available() else "cpu")

    # Use a predefined architecture
    if model_yaml in ["xvector", "rawnet2", "resnet34", "fastresnet34"]:

        if model_name is None:
            model = Xtractor(speaker_number, model_yaml)

        else:
            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

        model_archi = model_yaml

    # Here use a config file to build the architecture
    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)

    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"])

    """
    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 
    """
    training_df, validation_df = train_test_split(df, test_size=dataset_params["validation_ratio"], stratify=df["speaker_idx"])

    torch.manual_seed(dataset_params['seed'])

    training_set = SideSet(dataset_yaml,
                           set_type="train",
                           chunk_per_segment=-1,
                           overlap=dataset_params['train']['overlap'],
                           dataset_df=training_df,
                           output_format="pytorch",
                           )

    validation_set = SideSet(dataset_yaml,
                             set_type="validation",
                             dataset_df=validation_df,
                             output_format="pytorch")

    side_sampler = SideSampler(training_set.sessions['speaker_idx'],
                               speaker_number,
                               1,
                               100,
                               dataset_params["batch_size"])

    training_loader = DataLoader(training_set,
                                 batch_size=dataset_params["batch_size"],
                                 shuffle=False,
                                 drop_last=True,
                                 pin_memory=True,
                                 sampler=side_sampler,
                                 num_workers=num_thread,
                                 persistent_workers=True)

    validation_loader = DataLoader(validation_set,
                                   batch_size=dataset_params["batch_size"],
                                   drop_last=False,
                                   pin_memory=True,
                                   num_workers=num_thread,
                                   persistent_workers=False)

    """
    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}

    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)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=10 * training_loader.__len__(),
                                                gamma=0.95)