Discriminative Training of VBx Diarization
#!/usr/bin/env python3

"""
@Authors: Dominik Klement, Mireia Diez, Federico Landini (xkleme15@vutbr.cz, mireia@fit.vutbr.cz, landini@fit.vutbr.cz),
Brno University of Technology
"""

from VBx.diarization_lib import l2_norm, cos_similarity, twoGMMcalib_lin
from kaldi_io import read_vec_flt_ark
from itertools import groupby
import h5py
import numpy as np
from torch import from_numpy as torch_from_numpy
import fastcluster
from scipy.cluster.hierarchy import fcluster
from scipy.spatial.distance import squareform


class DataFileLoader:
    def __init__(self,
                 plda_mu,
                 segments_dir_path,
                 init_labels_dir_path,
                 xvec_ark_dir_path,
                 gt_labels_dir_path,
                 xvec_transform_path,
                 init_smoothing,
                 gt_label_type='PROBABILITIES',
                 ahc_threshold=0.015,
                 eval_mode=False):
        self.plda_mu = plda_mu
        self.segments_dir_path = segments_dir_path
        self.init_labels_dir_path = init_labels_dir_path
        self.xvec_ark_dir_path = xvec_ark_dir_path
        self.gt_labels_dir_path = gt_labels_dir_path
        self.xvec_transform_path = xvec_transform_path
        self.init_smoothing = init_smoothing
        self.gt_label_type = gt_label_type
        self.ahc_threshold = ahc_threshold

        # If true, we only need the xvectors and segments, and init labels, no gt labels
        self.eval_mode = eval_mode

        with h5py.File(xvec_transform_path, 'r') as f:
            self.mean1 = np.array(f['mean1'])
            self.mean2 = np.array(f['mean2'])
            self.lda = np.array(f['lda'])

    def load_file(self, file_name):
        arkit = read_vec_flt_ark(self.xvec_ark_dir_path + file_name + ".ark")
        recit = groupby(arkit, lambda e: e[0].rsplit('_', 1)[0])

        fn, segs = next(iter(recit))
        seg_names, xvecs = zip(*segs)
        x = np.array(xvecs)
        x = l2_norm(self.lda.T.dot((l2_norm(x - self.mean1)).transpose()).transpose() - self.mean2)

        if self.init_labels_dir_path is None:
            # Kaldi-like AHC of x-vectors (scr_mx is matrix of pairwise
            # similarities between all x-vectors)
            scr_mx = cos_similarity(x)
            # Figure out utterance specific args.threshold for AHC
            thr, _ = twoGMMcalib_lin(scr_mx.ravel())
            # output "labels" is an integer vector of speaker (cluster) ids
            scr_mx = squareform(-scr_mx, checks=False)
            lin_mat = fastcluster.linkage(
                scr_mx, method='average', preserve_input='False')
            del scr_mx
            adjust = abs(lin_mat[:, 2].min())
            lin_mat[:, 2] += adjust
            labels1st = fcluster(lin_mat, -(thr + self.ahc_threshold) + adjust, criterion='distance') - 1
        else:
            labels1st = np.loadtxt(self.init_labels_dir_path + file_name + ".init_labs", dtype=int)

        if self.eval_mode:
            max_label = np.max(labels1st)
            y_labels = None
        else:
            y_labels_N = np.loadtxt(self.gt_labels_dir_path + file_name + ".init_labs")
            if len(y_labels_N.shape) == 1:
                y_labels_N = np.expand_dims(y_labels_N, 1)

            max_label = np.maximum(np.shape(y_labels_N)[1] - 1, np.max(labels1st)).astype(int)

            if self.gt_label_type == 'SINGLE_SPK':
                """
                The i-th frame label is a one-hot encoding of the prevalent speaker in the frame 
                (i.e. speaking for most of the time).
                """
                y_labels_N = np.argmax(y_labels_N, axis=1)
            elif self.gt_label_type == 'ZERO_ONES':
                """
                The i-th frame label is the binary encoding of the presence of the speakers in the given frame
                (i.e. if a speaker s speaks in the given frame, the corresponding label is equal to 1).
                """
                y_labels_N = (y_labels_N > 0).astype(int)
            elif self.gt_label_type == 'PROBABILITIES':
                """
                The i-th frame label is proportional to the amount of speaker's speech in the frame
                (e.g. we have a 1.5s frame and 1s speaks 1st speaker and 0.5s speaks 2nd one, the corresponding labels
                are: (2/3, 1/3)).
                """
                y_labels = y_labels_N
            else:
                raise Exception("Ground truth label type not implemented")

            y_labels = np.zeros((len(y_labels_N), max_label + 1))
            if self.gt_label_type == 'SINGLE_SPK':
                y_labels[range(len(y_labels_N)), y_labels_N] = 1.0
            else:
                y_labels[:y_labels_N.shape[0], :y_labels_N.shape[1]] = y_labels_N

        qinit = np.zeros((len(labels1st), max_label + 1))
        qinit[range(len(labels1st)), labels1st] = 1.0
        fea = (x - self.plda_mu)

        return (
            self._read_xvector_timing_dict(self.segments_dir_path + file_name),
            seg_names,
            file_name,
            torch_from_numpy(fea),
            torch_from_numpy(qinit),
            np.empty(shape=(1, 1)) if self.eval_mode else torch_from_numpy(y_labels)
        )

    @staticmethod
    def _read_xvector_timing_dict(kaldi_segments):
        """ Loads kaldi 'segments' file with the timing information for individual x-vectors.
        Each line of the 'segments' file is expected to contain the following fields:
        - x-vector name (which needs to much the names of x-vectors loaded from kaldi archive)
        - file name of the recording from which the xvector is extracted
        - start time
        - end time
        Input:
            kaldi_segments - path (including file name) to the Kaldi 'segments' file
        Outputs:
             segs_dict[recording_file_name] = (array_of_xvector_names, array_of_start_and_end_times)
        """
        segs = np.loadtxt(kaldi_segments, dtype=object)
        split_by_filename = np.nonzero(segs[1:, 1] != segs[:-1, 1])[0] + 1
        return {s[0, 1]: (s[:, 0], s[:, 2:].astype(float)) for s in np.split(segs, split_by_filename)}

VBx/dataset/data_file_loader.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from torch.utils.data import DataLoader


def create_data_loader(dataset, batch_size=16, shuffle=True, num_workers=0):
    return DataLoader(dataset,
                      batch_size=batch_size,
                      shuffle=shuffle,
                      num_workers=num_workers,
                      collate_fn=lambda x: x)

VBx/dataset/data_loader.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from torch.utils.data import Dataset
from VBx.dataset.data_file_loader import DataFileLoader


class BasicDataset(Dataset):
    def __init__(self, file_list, file_loader: DataFileLoader):
        self.file_list = file_list
        self.file_loader = file_loader

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        if index > len(self.file_list):
            raise IndexError("Index out of range")

        return self.file_loader.load_file(self.file_list[index])

VBx/dataset/dataset.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from enum import Enum


class AvailableBackpropTypes(Enum):
    AFTER_CONVERGENCE = 1
    AFTER_EACH_ITER = 2

VBx/enums/backprop_types.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from enum import Enum


class GammaTypes(Enum):
    PROB = 1,
    LOG_PROB = 2,
    LOGIT = 3

VBx/enums/gamma_types.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from enum import Enum


class AvailableLosses(Enum):
    BCE = 1
    EDE = 2

VBx/enums/losses.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from scipy.optimize import linear_sum_assignment
import torch
import torch.nn.functional as F
from torch.nn.functional import logsigmoid
import numpy as np

from VBx.losses.loss import Loss
from VBx.enums.gamma_types import GammaTypes


class BCELoss(Loss):
    def __init__(self, use_scale=False, log_loss_scale=None, use_logprobs=True, use_adaptive_scale=False,
                 max_gamma_val=None):
        super().__init__(use_scale, log_loss_scale, use_logprobs, use_adaptive_scale, max_gamma_val,
                         need_gamma_type=GammaTypes.LOGIT)

    def __call__(self, decisions: torch.Tensor, target: torch.Tensor):
        """
        Inputs:
            decisions - TxS tensor (T - time, S - number of speakers) representing decisions made by a model
            target    - TxS tensor (T - time, S - number of speakers) representing ground truth
        """
        logits = super()._process_decisions(decisions)
        logits_t = logits.detach().transpose(0, 1)

        cost_mx = -logsigmoid(logits_t.unsqueeze(0)).bmm(target.unsqueeze(0)) - \
                  logsigmoid(-logits_t.unsqueeze(0)).bmm(1 - target.unsqueeze(0))
        pred_alig, ref_alig = linear_sum_assignment(cost_mx[0].to("cpu"))
        assert (np.all(pred_alig == np.arange(decisions.shape[-1])))
        target = target[:, ref_alig]

        # ToDo: Targets may have -1 values in the future if the conversations get padded for another batching approach.
        activation_loss = F.binary_cross_entropy_with_logits(logits, target, reduction='none')
        return torch.mean(activation_loss)

VBx/losses/bce.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from scipy.optimize import linear_sum_assignment
import torch

from VBx.losses.loss import Loss
from VBx.enums.gamma_types import GammaTypes


class EDELoss(Loss):
    """
    This loss is calculated as (fa + miss) / recording_length, which does not represent the DER metric as we know it,
    but rather Expected Detection Error Rate.
    """

    def __init__(self, use_scale=False, log_loss_scale=None, use_logprobs=True, use_adaptive_scale=False,
                 max_gamma_val=None):
        super().__init__(use_scale, log_loss_scale, use_logprobs, use_adaptive_scale, max_gamma_val,
                         need_gamma_type=GammaTypes.PROB)

    def __call__(self, decisions: torch.Tensor, target: torch.Tensor):
        """
        Inputs:
            decisions - TxS tensor (T - time, S - number of speakers) representing decisions made by a model
            target    - TxS tensor (T - time, S - number of speakers) representing ground truth
        """
        # Processed VBx responsibilities (probabilities).
        decisions = super()._process_decisions(decisions)
        decisions_t = torch.transpose(decisions, 0, 1)

        cost_mx2 = torch.matmul(1 - decisions_t, target) + torch.matmul(decisions_t, 1 - target)
        pred_alig, ref_alig = linear_sum_assignment(cost_mx2.detach())

        return torch.sum(cost_mx2[(pred_alig, ref_alig)]) / target.size()[0]

VBx/losses/ede.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from torch import log as torch_log, exp as torch_exp
from torch.nn.functional import softmax, log_softmax
import torch

from VBx.enums.gamma_types import GammaTypes
from VBx.utils.utils import logprobs_to_logodds


class Loss:
    def __init__(self, use_scale=False, log_loss_scale=None, use_logprobs=True, use_adaptive_scale=False,
                 max_gamma_val=0.95, need_gamma_type=GammaTypes.PROB):
        self.use_scale = use_scale
        self.use_logprobs = use_logprobs
        self.use_adaptive_scale = use_adaptive_scale
        self.max_gamma_val = max_gamma_val
        self.need_gamma_type = need_gamma_type

        if use_scale and log_loss_scale is None:
            raise 'loss_scale parameter cannot be None if use_scale is true.'

        self.log_loss_scale = log_loss_scale

    # ToDo: Merge scale_decisions and process_decisions because we may produce numerical imprecisions in some
    #  settings of "use_logits" and "need_logits" parameters.
    def _scale_decisions(self, decisions):
        # if use_logprobs is true then decisions are log responsibilities
        logprobs = decisions if self.use_logprobs else torch_log(decisions)
        scaled_logprobs = logprobs * torch.exp(self.log_loss_scale)

        if self.use_logprobs:
            return log_softmax(scaled_logprobs, dim=1)

        return softmax(scaled_logprobs, dim=1)

    def _process_decisions(self, decisions):
        if self.use_scale:
            decisions = self._scale_decisions(decisions)

        if self.need_gamma_type == GammaTypes.PROB:
            # The loss needs probabilities but decisions are log probabilities.
            if self.use_logprobs:
                decisions = torch_exp(decisions)
        elif self.need_gamma_type == GammaTypes.LOG_PROB:
            # The loss needs log probabilities but decisions are probabilities.
            if not self.use_logprobs:
                decisions = torch_log(decisions)
        else:
            # gamma_type = LOGIT
            decisions = logprobs_to_logodds(decisions if self.use_logprobs else torch_log(decisions))

        return decisions

VBx/losses/loss.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python

# @Authors: Lukas Burget, Mireia Diez, Dominik Klement (burget@fit.vutbr.cz, mireia@fit.vutbr.cz, xkleme15@vutbr.cz)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Revision History
#   L. Burget   20/1/2021 1:00AM - original version derived from the more
#                                  complex VB_diarization.py avaiable at
# https://github.com/BUTSpeechFIT/VBx/blob/e39af548bb41143a7136d08310765746192e34da/VBx/VB_diarization.py
#

import torch
from matplotlib import pyplot as plt
from torch import nn
import numpy as np
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter

from VBx.enums.losses import AvailableLosses
from VBx.enums.backprop_types import AvailableBackpropTypes
from VBx.losses.bce import BCELoss
from VBx.losses.ede import EDELoss


class DiscriminativeVB(nn.Module):
    def __init__(self,
                 device: torch.device,
                 config,
                 phi: torch.double,
                 fa: torch.double,
                 fb: torch.double,
                 loop_p: torch.double,
                 tr: torch.double,
                 init_smoothing: float,
                 epsilon: torch.double = 1e-4,
                 alpha_qinit: torch.double = 1.0,
                 alpha: torch.double = None,
                 invL: torch.double = None,
                 use_gmm: bool = False) -> None:

        super(DiscriminativeVB, self).__init__()

        self.device = device
        self._config = config
        self.use_gmm = use_gmm

        lp = loop_p * torch.ones(1, 1)
        lp_logit = torch.log(lp / (1 - lp))
        self.loop_prob_logit = nn.Parameter(lp_logit * torch.ones((1, 1), dtype=torch.double),
                                            requires_grad='lp' in self._config.trainable_params and \
                                                          not self.use_gmm and loop_p > 0)
        self.fa = nn.Parameter(fa * torch.ones((1, 1), dtype=torch.double),
                               requires_grad='fa' in self._config.trainable_params)
        self.fb = nn.Parameter(fb * torch.ones((1, 1), dtype=torch.double),
                               requires_grad='fb' in self._config.trainable_params)

        self.log_ism = nn.Parameter(torch.log(init_smoothing * torch.ones((1, 1), dtype=torch.double)),
                                    requires_grad='ism' in self._config.trainable_params)

        # Diagonal matrix of eigenvalues
        # self.phi = nn.Parameter(phi, requires_grad='phi' in self._config.trainable_params)
        self.original_log_phi = torch.clone(torch.log(phi)).detach()
        self.log_phi = nn.Parameter(torch.log(phi), requires_grad='phi' in self._config.trainable_params)

        # Transition matrix
        self.original_tr = torch.clone(tr).detach()
        self.tr = nn.Parameter(tr, requires_grad='tr' in self._config.trainable_params)

        train_loss_weights = self._config.use_loss_weights and \
                             self._config.backprop_type == AvailableBackpropTypes.AFTER_EACH_ITER and \
                             'loss_weights' in self._config.trainable_params

        if self._config.use_sigmoid_loss_weights:
            loss_weights_tensor = torch.zeros(
                (self._config.max_iters if self._config.avg_last_n_iters == -1 else self._config.avg_last_n_iters,),
                dtype=torch.double
            )
        else:
            loss_weights_tensor = torch.ones(
                (self._config.max_iters if self._config.avg_last_n_iters == -1 else self._config.avg_last_n_iters,),
                dtype=torch.double
            )
        self.loss_weights = nn.Parameter(loss_weights_tensor, requires_grad=train_loss_weights)

        train_loss_scale = 'loss_scale' in self._config.trainable_params and \
                           self._config.use_loss_scale and not self._config.use_adaptive_scale
        self.log_loss_scale = nn.Parameter(
            torch.log(torch.ones((1, 1), dtype=torch.double) * self._config.initial_loss_scale),
            requires_grad=train_loss_scale)

        self.max_iters = self._config.max_iters
        self.epsilon = epsilon
        self.alpha_qinit = alpha_qinit

        # Buffer that stores log_gamma values throughout the VB iterations
        self.log_gamma_buffer = []

        if alpha is not None:
            self.alpha = torch.tensor(alpha.copy(), dtype=torch.double, requires_grad=False)
        else:
            self.alpha = None

        if invL is not None:
            self.invL = torch.tensor(invL.copy(), dtype=torch.double, requires_grad=False)
        else:
            self.invL = None

        # Normal loss includes scaling and other postprocessing
        self._loss = self._setup_loss()

        # Raw loss is computed using the output gammas without any additional postprocessing
        self._raw_loss = self._setup_raw_loss()
        self._optimizer = self._setup_optimizer()

    def forward(self, X, pi, gamma):  # -> Tuple[torch.Tensor, torch.Tensor, List]
        self.log_gamma_buffer = []
        return self.VBx(X, pi, torch.nn.functional.softmax(torch.exp(self.log_ism) * gamma, dim=1))

    def compute_KL(self):
        Ew = torch.linalg.inv(torch.mm(self.original_tr.T, self.original_tr)).detach()
        Eb = torch.linalg.inv(
            torch.mm(self.original_tr.T / torch.exp(self.original_log_phi), self.original_tr)).detach()
        k = Ew.shape[0]

        tr_t_tr = torch.mm(self.tr.T, self.tr)
        KL_Ew = 0.5 * (-torch.logdet(torch.mm(Ew, tr_t_tr)) + torch.trace(torch.mm(Ew, tr_t_tr)) - k)
        KL_Eb = 0.5 * (torch.sum(self.log_phi) - torch.logdet(torch.mm(Eb, tr_t_tr)) + torch.trace(
            torch.mm(Eb, torch.mm(self.tr.T / torch.exp(self.log_phi), self.tr))) - k)

        return KL_Ew, KL_Eb

    def compute_loss(self, y_labels):
        if not self.log_gamma_buffer:
            raise 'Forward pass needs to be executed before computing a loss.'

        if not self.training:
            current_loss = self._loss(self.log_gamma_buffer[-1], y_labels)
            current_raw_loss = self._raw_loss(self.log_gamma_buffer[-1].detach(), y_labels)
            return current_loss.detach(), current_raw_loss.detach()

        if self._config.backprop_type == AvailableBackpropTypes.AFTER_CONVERGENCE:
            current_loss = self._loss(self.log_gamma_buffer[-1], y_labels)
            current_raw_loss = self._raw_loss(self.log_gamma_buffer[-1].detach(), y_labels)
        elif self._config.backprop_type == AvailableBackpropTypes.AFTER_EACH_ITER:
            losses = torch.stack(
                [self._loss(x, y_labels) for x in self.log_gamma_buffer[-self._config.avg_last_n_iters:]])
            raw_losses = torch.stack(
                [self._raw_loss(x.detach(), y_labels) for x in self.log_gamma_buffer[-self._config.avg_last_n_iters:]])

            if self._config.use_loss_weights:
                if self._config.use_sigmoid_loss_weights:
                    sgm_weights = F.sigmoid(self.loss_weights)
                    sgm_weights_2 = [(1 - torch.prod(sgm_weights[:i], dtype=torch.double)) for i in
                                     range(1, len(self.loss_weights) + 1)]
                    sws = sum(sgm_weights_2)
                    for i in range(len(losses)):
                        losses[i] *= sgm_weights_2[i] / sws
                else:
                    losses *= F.softmax(self.loss_weights, dim=0)

            current_loss = torch.sum(losses)
            current_raw_loss = torch.sum(raw_losses)

            if not self._config.use_loss_weights:
                current_loss /= losses.size()[0]
                current_raw_loss /= raw_losses.size()[0]
        else:
            raise 'Backprop type not implemented.'

        if self._config.use_regularization:
            # KL regularization
            KL_Ew, KL_Eb = self.compute_KL()
            current_loss += self._config.regularization_coeff_ew * KL_Ew + self._config.regularization_coeff_eb * KL_Eb

        return current_loss, current_raw_loss.detach()

    def process_gammas(self, log_gammas):
        """
        Processes the gammas the same way as loss does.
        """
        if self._config.loss == AvailableLosses.EDE:
            return self._loss._process_decisions(log_gammas)

        return torch.sigmoid(self._loss._process_decisions(log_gammas))

    def optimize_params(self, current_loss):
        self._optimizer.zero_grad()
        current_loss.backward()
        self._optimizer.step()

    def tb_log_params(self, tb_writer: SummaryWriter, tag, suffix, step):
        tb_writer.add_scalar(f'{tag}/fa_{suffix}', self.fa, step)
        tb_writer.add_scalar(f'{tag}/fb_{suffix}', self.fb, step)
        tb_writer.add_scalar(f'{tag}/pl_{suffix}', torch.sigmoid(self.loop_prob_logit), step)
        tb_writer.add_scalar(f'{tag}/loss_scale_{suffix}', torch.exp(self.log_loss_scale), step)
        tb_writer.add_histogram(f'{tag}/log_phi_{suffix}', self.log_phi, step)
        tb_writer.add_histogram(f'{tag}/tr_{suffix}', self.tr, step)

        tb_writer.add_scalar(f'{tag}/ism_{suffix}', torch.exp(self.log_ism), step)

        plt.clf()
        if self._config.use_sigmoid_loss_weights:
            sgm_weights = F.sigmoid(self.loss_weights.detach())
            sgm_weights_2 = [torch.prod(sgm_weights[:i]) for i in range(1, len(self.loss_weights) + 1)]
            plt.bar(x=range(len(self.loss_weights)), height=(1 - np.array(sgm_weights_2)))
        else:
            plt.bar(x=range(len(self.loss_weights)),
                    height=torch.nn.functional.softmax(self.loss_weights.detach(), dim=0).numpy())
        tb_writer.add_figure(f'{tag}/loss_weights_{suffix}', plt.gcf(), step)

    def tb_log_params_grads(self, tb_writer, tag, suffix, step):
        if self.fa.requires_grad:
            tb_writer.add_scalar(f'{tag}/fa_{suffix}', self.fa.grad, step)
        if self.fb.requires_grad:
            tb_writer.add_scalar(f'{tag}/fb_{suffix}', self.fb.grad, step)
        if self.loop_prob_logit.requires_grad:
            tb_writer.add_scalar(f'{tag}/pl_logit_{suffix}', self.loop_prob_logit.grad, step)
        if self.log_phi.requires_grad:
            tb_writer.add_scalar(f'{tag}/log_phi_l2_{suffix}', torch.norm(self.log_phi.grad), step)
            tb_writer.add_histogram(f'{tag}/log_phi_{suffix}', self.log_phi.grad, step)
        if self.tr.requires_grad:
            tb_writer.add_scalar(f'{tag}/tr_l2_{suffix}', torch.norm(self.tr.grad), step)
            tb_writer.add_histogram(f'{tag}/tr_{suffix}', self.tr.grad, step)
        if self.log_loss_scale.requires_grad:
            tb_writer.add_scalar(f'{tag}/log_loss_scale_{suffix}', self.log_loss_scale.grad, step)
        if self.loss_weights.requires_grad:
            tb_writer.add_scalar(f'{tag}/loss_weights_l2_{suffix}', torch.norm(self.loss_weights.grad), step)

    def save_model(self, path, ep, step):
        return torch.save({
            'model_state_dict': self.state_dict(),
            'optimizer_state_dict': self._optimizer.state_dict(),
            'epoch': ep,
            'step': step
        }, path)

    def load_model(self, path):
        checkpoint = torch.load(path)
        if 'ism' in checkpoint['model_state_dict']:
            checkpoint['model_state_dict']['log_ism'] = torch.log(checkpoint['model_state_dict']['ism'])
            del checkpoint['model_state_dict']['ism']
        self.load_state_dict(checkpoint['model_state_dict'])
        # self._optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        return checkpoint['epoch'], checkpoint['step']

    def VBx(self, X_in, pi_in, gamma):
        """
        Inputs:
        X           - T x D array, where columns are D dimensional feature vectors
                      (e.g. x-vectors) for T frames
        Phi         - D array with across-class covariance matrix diagonal.
                      The model assumes zero mean, diagonal across-class and
                      identity within-class covariance matrix.
        loopProb    - Probability of not switching speakers between frames
        Fa          - Scale sufficient statiscits
        Fb          - Speaker regularization coefficient Fb controls the final number of speakers
        pi          - If integer value, it sets the maximum number of speakers
                      that can be found in the utterance.
                      If vector, it is the initialization for speaker priors (see Outputs: pi)
        gamma       - An initialization for the matrix of responsibilities (see Outputs: gamma)
        maxIters    - The maximum number of VB iterations
        epsilon     - Stop iterating, if the obj. fun. improvement is less than epsilon
        alphaQInit  - Dirichlet concentraion parameter for initializing gamma
        ref         - T dim. integer vector with per frame reference speaker IDs (0:maxSpeakers)
        plot        - If set to True, plot per-frame marginal speaker posteriors 'gamma'
        return_model- Return also speaker model parameter
        alpha, invL - If provided, these are speaker model parameters used in the first iteration

        Outputs:
        gamma       - S x T matrix of responsibilities (marginal posteriors)
                      attributing each frame to one of S possible speakers
                      (S is defined by input parameter pi)
        pi          - S dimensional column vector of ML learned speaker priors.
                      This allows us to estimate the number of speaker in the
                      utterance as the probabilities of the redundant speaker
                      converge to zero.
        Li          - Values of auxiliary function (and DER and frame cross-entropy
                      between gamma and reference, if 'ref' is provided) over iterations.
        alpha, invL - Speaker model parameters returned only if return_model=True

        Reference:
          Landini F., Profant J., Diez M., Burget L.: Bayesian HMM clustering of
          x-vector sequences (VBx) in speaker diarization: theory, implementation
          and analysis on standard tasks
        """
        """
        The comments in the code refers to the equations from the paper above. Also
        the names of variables try to be consistent with the symbols in the paper.
        """
        X = torch.mm(X_in, self.tr.T)
        D = X.shape[1]

        if type(pi_in) is int:
            pi = torch.ones(pi_in, dtype=torch.double, requires_grad=False) / pi_in
        else:
            pi = pi_in

        if gamma is None:
            # initialize gamma from flat Dirichlet prior with
            # concentration parameter alphaQInit
            gamma = torch.from_numpy(np.random.gamma(self.alpha_qinit, size=(X.shape[0], len(pi))))
            gamma = gamma / gamma.sum(1, keepdim=True)

        assert (gamma.size()[1] == len(pi) and gamma.size()[0] == X.shape[0])

        phi = torch.exp(self.log_phi)

        G = -0.5 * (torch.sum(X ** 2, dim=1, keepdim=True) + D * torch.log(
            2 * torch.tensor(np.pi)))  # per-frame constant term in (23)
        V = torch.sqrt(phi)  # between (5) and (6)
        rho = X * V  # (18)
        Li = []
        log_gamma = torch.log(gamma)

        # We don't need transition matrix for GMM.
        if not self.use_gmm:
            loop_prob = torch.sigmoid(self.loop_prob_logit)
            if self._config.use_full_tr:
                tr = torch.eye(len(pi)) * loop_prob
                tr_off_diag = ((torch.ones_like(tr) - torch.eye(len(pi))) * ((1 - loop_prob) / (len(pi) - 1)))
                tr += tr_off_diag
            else:
                tr = torch.eye(len(pi)) * loop_prob + (1 - loop_prob) * pi  # (1) transition probability matrix

        for ii in range(self.max_iters if self.training else self._config.eval_max_iters):
            self.invL = 1.0 / (1 + self.fa / self.fb * gamma.sum(dim=0, keepdim=True).T * phi)  # (17) for all speakers
            self.alpha = self.fa / self.fb * self.invL * torch.tensordot(gamma.T, rho, dims=1)  # (16) for all speakers
            log_p_ = self.fa * (torch.tensordot(rho, self.alpha.T, dims=1) - 0.5 * torch.tensordot(
                (self.invL + self.alpha ** 2), phi, dims=1) + G)  # (23) for all speakers

            if self.use_gmm:
                log_gamma, gamma, log_pX_, pi = self._compute_gmm(pi, log_p_)
            else:
                log_gamma, gamma, log_pX_, pi, tr = self._compute_hmm(pi, log_p_, tr)

            ELBO = (log_pX_ + self.fb * 0.5 * torch.sum(
                torch.log(self.invL) - self.invL - self.alpha ** 2 + 1)).detach()  # (25)
            Li.append([ELBO])
            self.log_gamma_buffer.append(log_gamma)

            if ii > 0 and ELBO - Li[-2][0] < self.epsilon:
                if ELBO - Li[-2][0] < 0:
                    print('WARNING: Value of auxiliary function has decreased!')
                if (self.training and self._config.early_stop_vb) or (
                        not self.training and self._config.eval_early_stop_vb):
                    break
        return log_gamma, gamma, pi, Li

    def _compute_hmm(self, pi, log_p_, tr):
        loop_prob = torch.sigmoid(self.loop_prob_logit)
        if not self._config.use_full_tr:
            tr = torch.eye(len(pi)) * loop_prob + (1 - loop_prob) * pi  # (1) transition probability matrix

        # (19) gamma, (20) logA, (21) logB, (22) log_pX_
        log_gamma, log_pX_, logA, logB = DiscriminativeVB._forward_backward(log_p_, tr, pi)
        gamma = torch.exp(log_gamma)

        if self._config.use_full_tr:
            pi = gamma[0]
            tr = tr * torch.sum(torch.exp(torch.unsqueeze(logA[:-1, :], 2) +
                                          torch.unsqueeze(logB[1:, :] + log_p_[1:, :], 1) - log_pX_), dim=0)
            tr /= torch.unsqueeze(torch.sum(tr, dim=1), dim=1)
        else:
            pi = (gamma[0] + (1 - loop_prob) * pi * torch.sum(torch.exp(
                torch.logsumexp(logA[:-1], dim=1, keepdim=True) + log_p_[1:] + logB[1:] - log_pX_),
                dim=0)).reshape(-1)  # (24)
            pi = pi / pi.sum()

        return log_gamma, gamma, log_pX_, pi, tr

    def _compute_gmm(self, pi, log_p_):
        eps = 1e-8
        lpi = torch.log(pi)

        log_p_x = torch.logsumexp(log_p_ + lpi, dim=-1)  # expected (through q(Y)) per-frame log likelihood
        log_pX_ = torch.sum(log_p_x, dim=0)  # expected (through q(Y)) recording log likelihood
        log_gamma = log_p_ + lpi - torch.unsqueeze(log_p_x, dim=1)
        gamma = torch.exp(log_gamma)

        if self._config.use_full_tr:
            raise Exception('Not Implemented.')
        else:
            # PI cannot be equal directly to 0, because log_pi would be -inf, and it would result in nan values.
            pi = torch.sum(gamma, dim=0) + eps
            pi /= torch.sum(pi, dim=0)

        return log_gamma, gamma, log_pX_, pi

    @staticmethod
    def _forward_backward(lls, tr, ip):
        """
        Inputs:
            lls - matrix of per-frame log HMM state output probabilities
            tr  - transition probability matrix
            ip  - vector of initial state probabilities (i.e. starting in the state)
        Outputs:
            pi  - matrix of per-frame state occupation posteriors
            tll - total (forward) log-likelihood
            lfw - log forward probabilities
            lfw - log backward probabilities
        """
        eps = 1e-8
        ltr = torch.log(tr)
        lfw = [torch.full_like(lls[0], -torch.inf) for _ in range(lls.size()[0])]
        lbw = [torch.full_like(lls[0], -torch.inf) for _ in range(lls.size()[0])]
        lfw[0] = lls[0] + torch.log(ip + eps)
        lbw[-1] = torch.zeros_like(lls[0])

        for ii in range(1, len(lls)):
            lfw[ii] = lls[ii] + torch.logsumexp(lfw[ii - 1] + ltr.T, dim=1)

        for ii in reversed(range(len(lls) - 1)):
            lbw[ii] = torch.logsumexp(ltr + lls[ii + 1] + lbw[ii + 1], dim=1)

        lfw = torch.stack(lfw, dim=0)
        lbw = torch.stack(lbw, dim=0)

        tll = torch.logsumexp(lfw[-1], dim=0)
        log_gamma = lfw + lbw - tll
        return log_gamma, tll, lfw, lbw

    def _setup_loss(self):
        if self._config.loss == AvailableLosses.BCE:
            return BCELoss(use_scale=self._config.use_loss_scale, log_loss_scale=self.log_loss_scale, use_logprobs=True)
        elif self._config.loss == AvailableLosses.EDE:
            return EDELoss(use_scale=self._config.use_loss_scale, log_loss_scale=self.log_loss_scale, use_logprobs=True)
        else:
            raise 'Loss not implemented'

    def _setup_raw_loss(self):
        if self._config.loss == AvailableLosses.BCE:
            return BCELoss(use_scale=False, log_loss_scale=1, use_logprobs=True)
        elif self._config.loss == AvailableLosses.EDE:
            return EDELoss(use_scale=False, log_loss_scale=1, use_logprobs=True)
        else:
            raise 'Loss not implemented'

    def _setup_optimizer(self):
        model_params = dict(self.named_parameters())
        settings_list = []
        lr_params_settings = self._config.lr.keys()
        all_lr = 1e-2
        if 'all' in self._config.lr:
            settings_list = self.parameters()
            all_lr = self._config.lr['all']
        else:
            set_parameters = set()
            for lr_setting in lr_params_settings:
                # Each key can contain multiple parameter names separated by commas.
                setting_param_names = lr_setting.replace(' ', '').split(',')
                setting_params = []
                for pname in setting_param_names:
                    if pname == 'all' or pname == 'other':
                        continue

                    # config.json contains lp as loop probability, but we internally represent it
                    # as a logit and then run sigmoid to obtain a valid probability.
                    if pname == 'lp':
                        pname = 'loop_prob_logit'

                    if pname == 'phi':
                        pname = 'log_phi'

                    # Save parameters that are already set to know what parameters get "other" lr.
                    set_parameters.add(pname)
                    # Parameters corresponding to the specified names.
                    setting_params.append(model_params[pname])

                # Torch format of param-specific lr settings.
                if setting_params:
                    settings_list.append({
                        'params': setting_params,  # List of parameters.
                        'lr': self._config.lr[lr_setting]  # Corresponding LR.
                    })

            if 'other' in self._config.lr:
                # Save named params to set and do difference to set left params.
                all_params = set(map(lambda x: x[0], self.named_parameters()))
                notset_params = all_params.difference(set_parameters)
                settings_list.append({
                    'params': [model_params[x] for x in notset_params],
                    'lr': self._config.lr['other']
                })

        optimizer = optim.Adam(settings_list, lr=all_lr)
        return optimizer

VBx/models/VBx.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
The following script extracts x-vectors from input recordings and generates the corresponding ground truth labels.

@Authors: Mireia Diez, Lukas Burget, Federico Landini, Jan Profant
@Emails: mireia@fit.vutbr.cz, burget@fit.vutbr.cz, landini@fit.vutbr.cz, jan.profant@phonexia.com
"""

import argparse
import logging
import os
import time

import kaldi_io
import numpy as np
import onnxruntime
import soundfile as sf
import torch.backends

import features
from models.resnet import *

from utils.rttm_utils import rttm_to_hard_labels
from diarization_lib import mkdir_p

torch.backends.cudnn.enabled = False

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


class Timer(object):
    def __init__(self, name=None):
        self.name = name

    def __enter__(self):
        self.tstart = time.time()
        if self.name:
            logger.info(f'Start: {self.name}: ')

    def __exit__(self, type, value, traceback):
        if self.name:
            logger.info(f'End:   {self.name}: Elapsed: {time.time() - self.tstart} seconds')
        else:
            logger.info(f'End:   {self.name}: ')


def initialize_gpus(args):
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus


def load_utt(ark, utt, position):
    with open(ark, 'rb') as f:
        f.seek(position - len(utt) - 1)
        ark_key = kaldi_io.read_key(f)
        assert ark_key == utt, f'Keys does not match: `{ark_key}` and `{utt}`.'
        mat = kaldi_io.read_mat(f)
        return mat


def write_txt_vectors(path, data_dict):
    """ Write vectors file in text format.

    Args:
        path (str): path to txt file
        data_dict: (Dict[np.array]): name to array mapping
    """
    with open(path, 'w') as f:
        for name in sorted(data_dict):
            f.write(f'{name}  [ {" ".join(str(x) for x in data_dict[name])} ]{os.linesep}')


def get_embedding(fea, model, label_name=None, input_name=None, backend='pytorch'):
    if backend == 'pytorch':
        data = torch.from_numpy(fea).to(device)
        data = data[None, :, :]
        data = torch.transpose(data, 1, 2)
        spk_embeds = model(data)
        return spk_embeds.data.cpu().numpy()[0]
    elif backend == 'onnx':
        return model.run([label_name],
                         {input_name: fea.astype(np.float32).transpose()
                         [np.newaxis, :, :]})[0].squeeze()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpus', type=str, default='', help='use gpus (passed to CUDA_VISIBLE_DEVICES)')
    parser.add_argument('--model', required=False, type=str, default=None, help='name of the model')
    parser.add_argument('--weights', required=True, type=str, default=None, help='path to pretrained model weights')
    parser.add_argument('--model-file', required=False, type=str, default=None, help='path to model file')
    parser.add_argument('--ndim', required=False, type=int, default=64, help='dimensionality of features')
    parser.add_argument('--embed-dim', required=False, type=int, default=256, help='dimensionality of the emb')
    parser.add_argument('--seg-len', required=False, type=int, default=144, help='segment length')
    parser.add_argument('--seg-jump', required=False, type=int, default=24, help='segment jump')
    parser.add_argument('--in-file-list', required=True, type=str, help='input list of files')
    parser.add_argument('--in-lab-dir', required=True, type=str, help='input directory with VAD labels')
    parser.add_argument('--in-wav-dir', required=True, type=str, help='input directory with wavs')
    parser.add_argument('--out-ark-fn', required=True, type=str, help='output embedding file')
    parser.add_argument('--out-seg-fn', required=True, type=str, help='output segments file')
    parser.add_argument('--backend', required=False, default='pytorch', choices=['pytorch', 'onnx'],
                        help='backend that is used for x-vector extraction')
    parser.add_argument('--in-rttm-dir', required=False, type=str,
                        help='Input directory with reference rttm files')
    parser.add_argument('--out-xvector-label-GT-dir', required=True, type=str,
                        help='Output directory to save x-vector ground truth labels, obtained from reference rttm files')

    args = parser.parse_args()

    seg_len = args.seg_len
    seg_jump = args.seg_jump

    device = ''
    if args.gpus != '':
        logger.info(f'Using GPU: {args.gpus}')

        # gpu configuration
        initialize_gpus(args)
        device = torch.device(device='cuda')
    else:
        device = torch.device(device='cpu')

    model, label_name, input_name = '', None, None

    if args.backend == 'pytorch':
        if args.model_file is not None:
            model = torch.load(args.model_file)
            model = model.to(device)
        elif args.model is not None and args.weights is not None:
            model = eval(args.model)(feat_dim=args.ndim, embed_dim=args.embed_dim)
            model = model.to(device)
            checkpoint = torch.load(args.weights, map_location=device)
            model.load_state_dict(checkpoint['state_dict'], strict=False)
            model.eval()
    elif args.backend == 'onnx':
        model = onnxruntime.InferenceSession(args.weights)
        input_name = model.get_inputs()[0].name
        label_name = model.get_outputs()[0].name

    else:
        raise ValueError('Wrong combination of --model/--weights/--model_file '
                         'parameters provided (or not provided at all)')

    file_names = np.atleast_1d(np.loadtxt(args.in_file_list, dtype=object))

    with torch.no_grad():
        with open(args.out_seg_fn, 'w') as seg_file:
            with open(args.out_ark_fn, 'wb') as ark_file:
                for fn in file_names:
                    with Timer(f'Processing file {fn}'):
                        signal, samplerate = sf.read(f'{os.path.join(args.in_wav_dir, fn)}.wav')
                        labs = np.atleast_2d((np.loadtxt(f'{os.path.join(args.in_lab_dir, fn)}.lab',
                                                         usecols=(0, 1)) * samplerate).astype(int))
                        if samplerate == 8000:
                            noverlap = 120
                            winlen = 200
                            window = features.povey_window(winlen)
                            fbank_mx = features.mel_fbank_mx(
                                winlen, samplerate, NUMCHANS=64, LOFREQ=20.0, HIFREQ=3700, htk_bug=False)
                        elif samplerate == 16000:
                            noverlap = 240
                            winlen = 400
                            window = features.povey_window(winlen)
                            fbank_mx = features.mel_fbank_mx(
                                winlen, samplerate, NUMCHANS=64, LOFREQ=20.0, HIFREQ=7600, htk_bug=False)
                        else:
                            raise ValueError(f'Only 8kHz and 16kHz are supported. Got {samplerate} instead.')

                        LC = 150
                        RC = 149

                        np.random.seed(3)  # for reproducibility
                        signal = features.add_dither((signal * 2 ** 15).astype(int))

                        # for the labels
                        y_labels_fr = rttm_to_hard_labels(args.in_rttm_dir + fn + ".rttm", precision=samplerate)
                        my_y_list = []

                        for segnum in range(len(labs)):
                            seg_lab = y_labels_fr[0][labs[segnum, 0]:labs[segnum, 1]]
                            seg = signal[labs[segnum, 0]:labs[segnum, 1]]
                            if seg.shape[0] > 0.01 * samplerate:  # process segment only if longer than 0.01s
                                # Mirror noverlap//2 initial and final samples
                                seg = np.r_[seg[noverlap // 2 - 1::-1],
                                seg, seg[-1:-winlen // 2 - 1:-1]]
                                seg_lab = np.r_[
                                    seg_lab[noverlap // 2 - 1::-1], seg_lab, seg_lab[-1:-winlen // 2 - 1:-1]]
                                fea = features.fbank_htk(seg, window, noverlap, fbank_mx,
                                                         USEPOWER=True, ZMEANSOURCE=True)
                                fea = features.cmvn_floating_kaldi(fea, LC, RC, norm_vars=False).astype(np.float32)

                                slen = len(fea)
                                start = -seg_jump

                                for start in range(0, slen - seg_len, seg_jump):
                                    data = fea[start:start + seg_len]
                                    xvector = get_embedding(
                                        data, model, label_name=label_name, input_name=input_name, backend=args.backend)

                                    key = f'{fn}_{segnum:04}-{start:08}-{(start + seg_len):08}'
                                    if np.isnan(xvector).any():
                                        logger.warning(f'NaN found, not processing: {key}{os.linesep}')
                                    else:
                                        seg_start = round(labs[segnum, 0] / float(samplerate) + start / 100.0, 3)
                                        seg_end = round(
                                            labs[segnum, 0] / float(samplerate) + start / 100.0 + seg_len / 100.0, 3
                                        )
                                        seg_file.write(f'{key} {fn} {seg_start} {seg_end}{os.linesep}')
                                        kaldi_io.write_vec_flt(ark_file, xvector, key=key)

                                        y_labels_xvec = seg_lab[start * (winlen - noverlap):start * (
                                                    winlen - noverlap) + seg_len * (winlen - noverlap) + noverlap]
                                        my_y_list.append(np.sum(y_labels_xvec, axis=0) / len(y_labels_xvec))

                                if slen - start - seg_jump >= 10:
                                    data = fea[start + seg_jump:slen]
                                    xvector = get_embedding(
                                        data, model, label_name=label_name, input_name=input_name, backend=args.backend)

                                    key = f'{fn}_{segnum:04}-{(start + seg_jump):08}-{slen:08}'

                                    if np.isnan(xvector).any():
                                        logger.warning(f'NaN found, not processing: {key}{os.linesep}')
                                    else:
                                        y_labels_xvec = seg_lab[(start + seg_jump) * (winlen - noverlap):]
                                        my_y_list.append(np.sum(y_labels_xvec, axis=0) / len(y_labels_xvec))
                                        seg_start = round(
                                            labs[segnum, 0] / float(samplerate) + (start + seg_jump) / 100.0, 3
                                        )
                                        seg_end = round(labs[segnum, 1] / float(samplerate), 3)
                                        seg_file.write(f'{key} {fn} {seg_start} {seg_end}{os.linesep}')
                                        kaldi_io.write_vec_flt(ark_file, xvector, key=key)

                    # Generate GT labels.
                    y_labels = np.array(my_y_list)
                    mkdir_p(args.out_xvector_label_GT_dir)
                    with open(os.path.join(args.out_xvector_label_GT_dir, f'{fn}.init_labs'), 'w') as fp:
                        np.savetxt(fp, y_labels)

VBx/predict_and_label.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from types import SimpleNamespace

from VBx.enums.losses import AvailableLosses
from VBx.enums.backprop_types import AvailableBackpropTypes


def process_lr(args):
    if type(args.lr) == SimpleNamespace:
        lr = vars(args.lr)
        for key in lr:
            lr[key] = float(lr[key])
        final_lr = lr
    elif type(args.lr) == int or type(args.lr) == float:
        if args.lr <= 0:
            raise 'Learning rate has to be a positive number.'

        final_lr = {
            'all': args.lr
        }
    elif type(args.lr) == str:
        try:
            parsed_lr = float(args.lr)
            if parsed_lr <= 0:
                raise 'Learning rate has to be a positive number.'
            final_lr = {
                'all': parsed_lr
            }
        except ValueError:
            raise 'Wrong "lr" value. It should be either a dictionary or a positive number.'

    else:
        raise 'Wrong "lr" value. It should be either a dictionary or a nonzero number.'

    return final_lr


# Validate arguments
def process_args(args):
    if args.avg_last_n_iters != -1 and args.avg_last_n_iters < 0:
        raise '\"avg_last_n_iters\" has to be -1 (averaging all VB iterations) or a positive number.'

    if args.backprop_type == 'after_each_iter':
        args.backprop_type = AvailableBackpropTypes.AFTER_EACH_ITER
    elif args.backprop_type == 'after_convergence':
        args.backprop_type = AvailableBackpropTypes.AFTER_CONVERGENCE
    else:
        raise 'Wrong backprop type. Use after_each_iter or after_convergence.'

    if args.batch_size < 1:
        raise 'Batch size has to be at least 1.'

    if args.eval_max_iters < 1:
        raise 'Number of eval VB iterations has to be at least 1.'

    if args.epochs < 1:
        raise "Number of epochs has to be at least 1."

    if args.initial_loss_scale < 0:
        raise "Initial loss scale has to be positive."

    if args.loss == 'BCE':
        args.loss = AvailableLosses.BCE
    elif args.loss == 'EDE':
        args.loss = AvailableLosses.EDE
    else:
        raise 'Loss not implemented. Use BCE or EDE.'

    # Convert learning rate namespace to dictionary
    args.lr = process_lr(args)

    if args.max_iters < 1:
        raise '"max_iters" should be at least 1.'

    if args.regularization_coeff_eb < 0:
        raise 'Regularization coefficient has to be a non-negative number.'

    if args.regularization_coeff_ew < 0:
        raise 'Regularization coefficient has to be a non-negative number.'

    args.trainable_params = set(args.trainable_params)

VBx/utils/config.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

from pyannote.core import Segment, Annotation
from pyannote.metrics.diarization import DiarizationErrorRate
from pyannote.database.loader import RTTMLoader
import torch
import numpy as np

class DER:
    def __init__(self, file_list, gt_rttm_path):
        self.rttms = dict()
        for f in file_list:
            self.rttms[f] = RTTMLoader(f'{gt_rttm_path}/{f}.rttm').loaded_[f]

    def der_tensor(self, start, end, decisions: torch.tensor, file_name):
        """
        Compute DER between the decisions and the reference torch tensors
        with dimension TxN, where T is the time axis and N is the max number of speakers.

        Inputs:
            start     - ...
            end       - ...
            decisions - TxN torch tensor containing binary decisions.
            reference - TxN torch tensor containing binary labels.
        Returns:
            Tuple containing the following elements:
            (DER score between the decisions and the reference labels,
             Speech amount in the given file).
        """

        annot: Annotation = DER._transfer_binary_decisions_to_pyannote_annotation(start, end, decisions)

        speech_amount = 0
        for seg, _ in self.rttms[file_name].itertracks():
            speech_amount += seg.duration

        return DiarizationErrorRate()(self.rttms[file_name], annot), speech_amount

    @staticmethod
    def _transfer_binary_decisions_to_pyannote_annotation(start, end, decisions) -> Annotation:
        starts, ends, out_labels = DER._merge_adjacent_labels(start, end, decisions)

        annotation_sys = Annotation()

        for s, e, l in zip(starts, ends, out_labels):
            s = Segment(s, e)
            annotation_sys[s, annotation_sys.new_track(s)] = l

        return annotation_sys

    @staticmethod
    def _merge_adjacent_labels(starts, ends, labels):
        """ Labeled segments defined as start and end times are compacted in such a way that
        adjacent or overlapping segments with the same label are merged. Overlapping
        segments with different labels are further adjusted not to overlap (the boundary
        is set in the middle of the original overlap).
        Input:
             starts - array of segment start times in seconds
             ends   - array of segment end times in seconds
             labels - array of segment labels (of any type)
        Outputs:
              starts, ends, labels - compacted and ajusted version of the input arrays
        """
        # Merge neighbouring (or overlaping) segments with the same label
        adjacent_or_overlap = np.logical_or(np.isclose(ends[:-1], starts[1:]), ends[:-1] > starts[1:])
        to_split = np.nonzero(np.logical_or(~adjacent_or_overlap, labels[1:] != labels[:-1]))[0]
        starts = starts[np.r_[0, to_split + 1]]
        ends = ends[np.r_[to_split, -1]]
        labels = labels[np.r_[0, to_split + 1]]

        # Fix starts and ends times for overlapping segments
        overlaping = np.nonzero(starts[1:] < ends[:-1])[0]
        ends[overlaping] = starts[overlaping + 1] = (ends[overlaping] + starts[overlaping + 1]) / 2.0
        return starts, ends, labels

VBx/utils/metrics.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python

# Copyright Brno University of Technology
# Licensed under the Apache License, Version 2.0 (the "License")

import numpy as np
import os
from collections import OrderedDict
from typing import List, Tuple


def hard_labels_to_rttm(
    matrix: np.ndarray,
    spk_labels: List[str],
    utt_id: str,
    rttm_path: str,
    precision: float
):
    """
        reads an NfxNs matrix containing hard speaker activities (labels 1/0)
        at the given precision. The speaker labels should correspond to the
        order in the array and are used to assign the speech segments in the
        rttm. The utterance ID is used for the second field of each rttm row.
    """
    matrix_extended = np.concatenate((
        np.zeros((1, matrix.shape[1])),
        matrix,
        np.zeros((1, matrix.shape[1]))), axis=0)
    changes_dict = {}
    for s in range(len(spk_labels)):
        spk_changes = np.where(matrix_extended[1:, s] - matrix_extended[:-1, s]
                               )[0].astype(float)
        if spk_changes.shape[0] > 0:
            if spk_changes[-1] == matrix.shape[0]:
                spk_changes[-1] -= 1  # avoid reading out of array
            beg = spk_changes[:-1]
            end = spk_changes[1:]
            # So far, beg and end include the silences in between
            beg = beg[::2]
            end = end[::2]
            assert beg.shape[0] == end.shape[0], "Amount of beginning and \
                                               end of segments do not match."
            for pos in range(beg.shape[0]):
                time_beg = beg[pos] / precision
                time_length = (end[pos] - beg[pos]) / precision
                changes_dict[(time_beg, s)] = f"SPEAKER {utt_id} 1 {time_beg} {time_length} <NA> <NA> {spk_labels[s]} <NA> <NA>\n"
    changes_dict = OrderedDict(sorted(changes_dict.items()))
    if not os.path.exists(os.path.dirname(rttm_path)):
        os.makedirs(os.path.dirname(rttm_path))
    with open(rttm_path, 'w') as f:
        for k, v in changes_dict.items():
            f.write(v)


def rttm_to_hard_labels(
    rttm_path: str,
    precision: float,
    length: float = -1
) -> Tuple[np.ndarray, List[str]]:
    """
        reads the rttm and returns a NfxNs matrix encoding the segments in
        which each speaker is present (labels 1/0) at the given precision.
        Ns is the number of speakers and Nf is the resulting number of frames,
        according to the parameters given.
        Nf might be shorter than the real length of the utterance, as final
        silence parts cannot be recovered from the rttm.
        If length is defined (s), it is to account for that extra silence.
        In case of silence all speakers are labeled with 0.
        In case of overlap all speakers involved are marked with 1.
        The function assumes that the rttm only contains speaker turns (no
        silence segments).
        The overlaps are extracted from the speaker turn collisions.
    """
    # each row is a turn, columns denote beginning (s) and duration (s) of turn
    data = np.loadtxt(rttm_path, usecols=[3, 4])
    # speaker id of each turn
    spks = np.loadtxt(rttm_path, usecols=[7], dtype='str')
    spk_ids = np.unique(spks)
    Ns = len(spk_ids)
    if data.shape[0] == 2 and len(data.shape) < 2:  # if only one segment
        data = np.asarray([data])
        spks = np.asarray([spks])
    # length of the file (s) that can be recovered from the rttm,
    # there might be extra silence at the end
    #len_file = data[-1][0]+data[-1][1] #last item in the rttm might not be the one containing the last segment of speech
    len_file = np.max(np.sum(data,axis=1)) 
    if length > len_file:
        len_file = length

    # matrix in given precision
    matrix = np.zeros([int(round(len_file*precision)), Ns])
    # ranges to mark each turn
    ranges = np.around((np.array([data[:, 0],
                        data[:, 0]+data[:, 1]]).T*precision)).astype(int)

    for s in range(Ns):  # loop over speakers
        # loop all the turns of the speaker
        for init_end in ranges[spks == spk_ids[s], :]:
            matrix[init_end[0]:init_end[1], s] = 1  # mark the spk
    return matrix, spk_ids

VBx/utils/rttm_utils.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python3

# @Authors: Lukas Burget (burget@fit.vutbr.cz), Brno University of Technology

import torch


def logprobs_to_logodds(logprobs):
    """ Transform a matrix of multi-class (unnormalized) log probabilities into (binary) log odds (logits).
    Input:  logprobs - TxC matrix, where T is number of examples and C is number of classes
    Output: logodds  - TxC matrix
    """
    C = logprobs.shape[1]
    logodds = torch.zeros_like(logprobs)

    for m in range(C):
        impostor_indices = torch.where(torch.arange(C) != m)
        logodds[:, m] = logprobs[:, m] - torch.logsumexp(logprobs.T[impostor_indices], dim=0)

    return logodds

VBx/utils/utils.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python

# @Authors: Lukas Burget, Mireia Diez, Federico Landini, Jan Profant
# @Emails: burget@fit.vutbr.cz, mireia@fit.vutbr.cz, landini@fit.vutbr.cz, jan.profant@phonexia.com
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# The recipe consists in doing Agglomerative Hierachical Clustering on
# x-vectors in a first step. 

import argparse
import os
import itertools

import fastcluster
import h5py
import kaldi_io
import numpy as np
from scipy.cluster.hierarchy import fcluster
from scipy.spatial.distance import squareform
from scipy.linalg import eigh
 
from VBx.diarization_lib import read_xvector_timing_dict, l2_norm, \
    cos_similarity, twoGMMcalib_lin, merge_adjacent_labels, mkdir_p
from VBx.utils.kaldi_utils import read_plda
import VBx.features as features
from VBx.utils.rttm_utils import rttm_to_hard_labels

def write_output(fp, out_labels, starts, ends):
    for label, seg_start, seg_end in zip(out_labels, starts, ends):
        fp.write(f'SPEAKER {file_name} 1 {seg_start:03f} {seg_end - seg_start:03f} '
                 f'<NA> <NA> {label + 1} <NA> <NA>{os.linesep}')

def write_xveclabs(fp, out_labels):
    for label in out_labels:
        fp.write(f'{label}{os.linesep}')

def mostFrequent(arr):
  maxcount = 0;
  element_having_max_freq = 0;
  farr = arr.flatten()
  for i in np.unique(farr):
    count = np.sum(farr==i)
    if(count > maxcount):
      maxcount = count
      element_having_max_freq = i
    
  return element_having_max_freq;


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--out-AHCinit-dir', required=True, type=str,
                        help='Output directory to save AHC initialization rttm and per x-vector label files')
    parser.add_argument('--xvec-ark-file', required=True, type=str,
                        help='Kaldi ark file with x-vectors from one or more '
                             'input recordings. Attention: all x-vectors from '
                             'one recording must be in one ark file')
    parser.add_argument('--segments-file', required=True, type=str,
                        help='File with x-vector timing info. See '
                        'diarization_lib.read_xvector_timing_dict')
    parser.add_argument('--xvec-transform', required=True, type=str,
                        help='path to x-vector transformation h5 file')
    parser.add_argument('--plda-file', required=True, type=str,
                        help='File with PLDA model in Kaldi format used for '
                        'AHC and VB-HMM x-vector clustering')
    parser.add_argument('--threshold', required=True, type=float,
                        help='args.threshold (bias) used for AHC')
    parser.add_argument('--target-energy', required=False, type=float,
                        default=1.0,
                        help='Parameter affecting AHC if the similarity '
                             'matrix is obtained with PLDA. See '
                             'diarization_lib.kaldi_ivector_plda_scoring_dense')
    parser.add_argument('--max-speakers', required=False, type=int, default=-1,
                        help='Imposes a constraint on max number of clusters found by AHC. -1 means no maximum.')

    args = parser.parse_args()

    # segments file with x-vector timing information
    segs_dict = read_xvector_timing_dict(args.segments_file)

    kaldi_plda = read_plda(args.plda_file)
    plda_mu, plda_tr, plda_psi = kaldi_plda
    W = np.linalg.inv(plda_tr.T.dot(plda_tr))
    B = np.linalg.inv((plda_tr.T / plda_psi).dot(plda_tr))
    acvar, wccn = eigh(B, W)
    plda_psi = acvar[::-1]
    plda_tr = wccn.T[::-1]

    # Open ark file with x-vectors and in each iteration of the following
    # for-loop read a batch of x-vectors corresponding to one recording
    arkit = kaldi_io.read_vec_flt_ark(args.xvec_ark_file)
    # group xvectors in ark by recording name
    recit = itertools.groupby(arkit, lambda e: e[0].rsplit('_', 1)[0])
    for file_name, segs in recit:
        print(file_name)
        seg_names, xvecs = zip(*segs)
        x = np.array(xvecs)

        with h5py.File(args.xvec_transform, 'r') as f:
            mean1 = np.array(f['mean1'])
            mean2 = np.array(f['mean2'])
            lda = np.array(f['lda'])
            x = l2_norm(lda.T.dot((l2_norm(x - mean1)).transpose()).transpose() - mean2)


        # Kaldi-like AHC of x-vectors (scr_mx is matrix of pairwise
        # similarities between all x-vectors)
        scr_mx = cos_similarity(x)
        # Figure out utterance specific args.threshold for AHC
        thr, _ = twoGMMcalib_lin(scr_mx.ravel())
        # output "labels" is an integer vector of speaker (cluster) ids
        scr_mx = squareform(-scr_mx, checks=False)
        lin_mat = fastcluster.linkage(
            scr_mx, method='average', preserve_input='False')
        del scr_mx
        adjust = abs(lin_mat[:, 2].min())
        lin_mat[:, 2] += adjust
        labels1st = fcluster(lin_mat, -(thr + args.threshold) + adjust,
            criterion='distance') - 1
        if args.max_speakers != -1 and np.unique(labels1st) > args.max_speakers:
            labels1st = fcluster(lin_mat, args.max_speakers, criterion='maxclust') - 1
 
        #save x-vector granularity AHC labels
        mkdir_p(args.out_AHCinit_dir)
        with open(os.path.join(args.out_AHCinit_dir, f'{file_name}.init_labs'), 'w') as fp:
            write_xveclabs(fp, labels1st)

        assert(np.all(segs_dict[file_name][0] == np.array(seg_names)))
        start, end = segs_dict[file_name][1].T

        #save standard rttm files of the AHC initialization
        starts, ends, out_labels = merge_adjacent_labels(start, end, labels1st)
        with open(os.path.join(args.out_AHCinit_dir, f'{file_name}.rttm'), 'w') as fp:
            write_output(fp, out_labels, starts, ends)



initAHC.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

import argparse
import os
import numpy as np
from scipy.linalg import eigh
import torch
import torch.distributed as dist

from VBx.diarization_lib import merge_adjacent_labels, mkdir_p
from VBx.utils.kaldi_utils import read_plda
from VBx.models import VBx as VBx
from VBx.dataset.dataset import BasicDataset
from VBx.dataset.data_file_loader import DataFileLoader
from VBx.utils.config import process_args
from VBx.utils.metrics import DER


def read_kaldi_plda(path):
    plda_mu, plda_tr, plda_psi = read_plda(path)
    W = np.linalg.inv(plda_tr.T.dot(plda_tr))
    B = np.linalg.inv((plda_tr.T / plda_psi).dot(plda_tr))
    acvar, wccn = eigh(B, W)
    plda_psi = acvar[::-1]
    plda_tr = wccn.T[::-1]
    return plda_mu, plda_psi, plda_tr


def trace_handler(prof):
    print(prof.key_averages().table(row_limit=-1))
    prof.export_chrome_trace("test_trace_" + str(prof.step_num) + ".json")


def eval_model(model, dset, run_dist, der_wrapper: DER, rank, world_size):
    model.eval()
    with torch.no_grad():
        all_ders = []

        for batch in dset:
            batch_der = []
            for sd, sn, fn, x, qinit, y in batch:
                pi_init = torch.ones(qinit.size()[1],
                                     dtype=torch.double,
                                     requires_grad=False) / qinit.size()[1]

                log_gamma, gamma, pi, Li = model(x, pi=pi_init, gamma=qinit)

                assert (np.all(sd[fn][0] == np.array(sn)))
                start, end = sd[fn][1].T

                labels1st = torch.argmax(gamma, dim=1).numpy()
                der, speech_amount = der_wrapper.der_tensor(start, end, labels1st, fn)
                batch_der.append((fn, der, speech_amount))

            if run_dist:
                batch_ders = [[] for _ in range(world_size)]
                dist.all_gather_object(batch_ders, batch_der)
            else:
                batch_ders = [batch_der]

            if rank == 0:
                all_ders.append(batch_ders)

        if rank == 0:
            # Get rid of repetitive files (sampling w repetitions is used for dist. batching)
            files_dict = dict()
            for batch_ders in all_ders:
                for batch_der in batch_ders:
                    for fn, der, speech_amount in batch_der:
                        files_dict[fn] = (der * speech_amount, speech_amount)

            total_der = sum([files_dict[x][0] for x in files_dict])
            total_speech = sum([files_dict[x][1] for x in files_dict])
            return 100.0 * total_der / total_speech


def write_output(fp, file_name, out_labels, starts, ends):
    for label, seg_start, seg_end in zip(out_labels, starts, ends):
        fp.write(f'SPEAKER {file_name} 1 {seg_start:03f} {seg_end - seg_start:03f} '
                 f'<NA> <NA> {label + 1} <NA> <NA>{os.linesep}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--in-file-list', required=False, type=str,
                        help='txt list of files')
    parser.add_argument('--out-rttm-dir', required=True, type=str,
                        help='Directory to store output rttm files')
    parser.add_argument('--xvec-ark-dir', required=True, type=str,
                        help='Kaldi ark file with x-vectors from one or more '
                             'input recordings. Attention: all x-vectors from '
                             'one recording must be in one ark file')
    parser.add_argument('--segments-dir', required=True, type=str,
                        help='File with x-vector timing info. See '
                             'diarization_lib.read_xvector_timing_dict')
    parser.add_argument('--xvec-transform', required=True, type=str,
                        help='path to x-vector transformation h5 file')
    parser.add_argument('--plda-file', required=True, type=str,
                        help='File with PLDA model in Kaldi format used for '
                             'AHC and VB-HMM x-vector clustering')
    parser.add_argument('--threshold', required=True, type=float,
                        help='args.threshold (bias) used for AHC')
    parser.add_argument('--lda-dim', required=True, type=int,
                        help='For VB-HMM, x-vectors are reduced to this '
                             'dimensionality using LDA')
    parser.add_argument('--target-energy', required=False, type=float,
                        default=1.0,
                        help='Parameter affecting AHC if the similarity '
                             'matrix is obtained with PLDA. See '
                             'diarization_lib.kaldi_ivector_plda_scoring_dense')
    parser.add_argument('--in-INITlabels-dir', required=False, type=str,
                        help='Directory with the AHC initialization xvector labels. '
                             'If not provided, the system will apply AHC first.')
    parser.add_argument('--model-path', required=True, type=str, help="Path to saved .pt model checkpoint.")
    parser.add_argument('--use-gmm', required=False, default=False, action='store_true',
                        help='If present, GMM is used instead of HMM, i.e. ploop=0.')
    parser.add_argument('--output-2nd', type=bool, default=False,
                        help='If present, system will also output the second most probable speaker.')

    # Default config. Those parameters don't have to be set.
    parser.add_argument('--avg-last-n-iters', type=int, default=-1,
                        help='Dictates how many last VB iterations are averaged to compute gradients '
                             '(in case of using after_each_iter backprop_type)')
    parser.add_argument('--backprop-type', type=str, default='after_each_iter',
                        help='Gradient computation type: after_convergence, or after_each_iter.')
    parser.add_argument('--batch-size', type=int, default=8, help='Batch size.')
    parser.add_argument('--early-stop-vb', type=bool, default=False,
                        help='If True, VB inference is stopped during training if ELBO stops increasing.')
    parser.add_argument('--eval-max-iters', type=int, default=40,
                        help='Max number of VB iterations used during inference/evaluation.')
    parser.add_argument('--eval-early-stop-vb', type=bool, default=True,
                        help='Same as --early-stop-vb but during evaluation.')
    parser.add_argument('--epochs', type=int, default=500, help='Number of passes through the whole training set.')
    parser.add_argument('--initial-loss-scale', type=float, default=1,
                        help='Initial value of loss scale (calibration).')
    parser.add_argument('--loss', type=str, default='EDE', help='Training loss, either BCE or DER.')
    parser.add_argument('--lr', default=1e-3, help='Learning rates.')
    parser.add_argument('--max-iters', type=int, default=10,
                        help='Max number of VB iterations during training (if early_stop_vb is False, than it\'')
    parser.add_argument('--regularization-coeff-eb', default=0, type=float,
                        help='KL divergence between-speaker covariance matrix regularization coefficient.')
    parser.add_argument('--regularization-coeff-ew', default=0, type=float,
                        help='KL divergence within-speaker covariance matrix regularization coefficient.')
    parser.add_argument('--trainable-params', default=['fa', 'fb', 'lp', 'ism'], nargs='*',
                        help='List of trainable parameters.')
    parser.add_argument('--use-full-tr', type=bool, default=False,
                        help='Use full transition matrix (i.e. allow HMM to have connection '
                             'between all pairs of nodes).')
    parser.add_argument('--use-loss-scale', type=bool, default=False,
                        help='If true, loss scale (calibration) is used.')
    parser.add_argument('--use-loss-weights', type=bool, default=False,
                        help='Gradients are computed as a weighted average of the VB iteration gradients.')
    parser.add_argument('--use-regularization', type=bool, default=False,
                        help='If True, KL divergence regularization towards the original (generatively trained) '
                             'PLDA covariance matrices is used.')
    parser.add_argument('--use-sigmoid-loss-weights', type=bool, default=False,
                        help='If true, loss weights for iteration i are computed as 1 - product of '
                             'weights for weights i,i+1, ..., n.')

    args = parser.parse_args()

    process_args(args)

    plda_mu, plda_psi, plda_tr = read_kaldi_plda(args.plda_file)

    model = VBx.DiscriminativeVB(device=torch.device("cpu"),
                                 phi=torch.from_numpy(plda_psi[:args.lda_dim].copy()),
                                 tr=torch.from_numpy(plda_tr.copy()),
                                 # fa,fb, loop_p, init_smoothing are loaded later on from .pt checkpoint.
                                 fa=1,
                                 fb=1,
                                 loop_p=0.5,
                                 config=args,
                                 init_smoothing=7,
                                 epsilon=1e-6,
                                 use_gmm=args.use_gmm)
    model.load_model(args.model_path)

    data_list = np.loadtxt(args.in_file_list, dtype=object)
    dfl = DataFileLoader(plda_mu,
                         args.segments_dir,
                         args.in_INITlabels_dir,
                         args.xvec_ark_dir,
                         None,
                         args.xvec_transform,
                         7,
                         ahc_threshold=args.threshold,
                         eval_mode=True)
    dset = BasicDataset(data_list, dfl)

    model.eval()
    with torch.no_grad():
        for segs_dict, seg_names, file_name, x, qinit, y in dset:
            pi_init = torch.ones(qinit.size()[1],
                                 dtype=torch.double,
                                 requires_grad=False) / qinit.size()[1]

            log_gamma, gamma, pi, Li = model(x, pi=pi_init, gamma=qinit)
            labels1st = torch.argsort(-gamma, dim=1)[:, 0].detach().numpy()
            labels2nd = np.argsort(-gamma, axis=1)[:, 1].detach().numpy()

            start, end = segs_dict[file_name][1].T

            starts, ends, out_labels = merge_adjacent_labels(start, end, labels1st)
            mkdir_p(args.out_rttm_dir)
            with open(os.path.join(args.out_rttm_dir, f'{file_name}.rttm'), 'w') as fp:
                write_output(fp, file_name, out_labels, starts, ends)

            if args.output_2nd and args.init.endswith('VB') and gamma.shape[1] > 1:
                starts, ends, out_labels2 = merge_adjacent_labels(start, end, labels2nd)
                output_rttm_dir = f'{args.out_rttm_dir}2nd'
                mkdir_p(output_rttm_dir)
                with open(os.path.join(output_rttm_dir, f'{file_name}.rttm'), 'w') as fp:
                    write_output(fp, file_name, out_labels2, starts, ends)

vbhmm_infer.py
BUTSpeechFIT/DVBx
#!/usr/bin/env python

# @Authors: Dominik Klement (xkleme15@vutbr.cz), Brno University of Technology

import os
from math import ceil
from shutil import copy as shell_copy
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from scipy.linalg import eigh
from scipy.optimize import linear_sum_assignment
import torch
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from socket import gethostname
from datetime import datetime
import yamlargparse

from VBx.dataset.data_loader import create_data_loader
from VBx.utils.kaldi_utils import read_plda
from VBx.models import VBx as VBx
from VBx.dataset.dataset import BasicDataset
from VBx.dataset.data_file_loader import DataFileLoader
from VBx.utils.config import process_args
from VBx.utils.metrics import DER

mpl.rcParams['figure.dpi'] = 150


def read_kaldi_plda(path):
    plda_mu, plda_tr, plda_psi = read_plda(path)
    W = np.linalg.inv(plda_tr.T.dot(plda_tr))
    B = np.linalg.inv((plda_tr.T / plda_psi).dot(plda_tr))
    acvar, wccn = eigh(B, W)
    plda_psi = acvar[::-1]
    plda_tr = wccn.T[::-1]
    return plda_mu, plda_psi, plda_tr


def trace_handler(prof):
    print(prof.key_averages().table(row_limit=-1))
    prof.export_chrome_trace("test_trace_" + str(prof.step_num) + ".json")


def plot_gammas(gammas, ys_alig, zero_ones_labels=False):
    colors = ['#61a4f1', '#a7f3e5', '#083167', '#c3a443', '#880b08', '#9f2903', '#7c3d2a',
              '#63d0fc',
              '#edce04', '#fa3dcf', '#82cefb', '#5f54b5', '#65388d', '#03a6d0', '#84cb7a',
              '#485a18',
              '#474c85', '#67edbb', '#85d7e6', '#7fb5ea', '#96591d', '#298b53', '#522f12',
              '#99c542',
              '#3b17f0', '#e4ef46', '#7d4247', '#9ded6f', '#260995', '#64ec6b', '#33f63d',
              '#7e5a90',
              '#448cc8', '#36b23f', '#eafb65', '#99f0e0', '#c059f3', '#02b8b5', '#c3e236',
              '#5f8ed3',
              "#000000", "#FFFF00", "#1CE6FF", "#FF34FF", "#FF4A46", "#008941", "#006FA6", "#A30059",
              "#FFDBE5", "#7A4900", "#0000A6", "#63FFAC", "#B79762", "#004D43", "#8FB0FF", "#997D87",
              "#5A0007", "#809693", "#FEFFE6", "#1B4400", "#4FC601", "#3B5DFF", "#4A3B53", "#FF2F80",
              "#61615A", "#BA0900", "#6B7900", "#00C2A0", "#FFAA92", "#FF90C9", "#B903AA", "#D16100",
              "#DDEFFF", "#000035", "#7B4F4B", "#A1C299", "#300018", "#0AA6D8", "#013349", "#00846F",
              "#372101", "#FFB500", "#C2FFED", "#A079BF", "#CC0744", "#C0B9B2", "#C2FF99", "#001E09",
              "#00489C", "#6F0062", "#0CBD66", "#EEC3FF", "#456D75", "#B77B68", "#7A87A1", "#788D66",
              "#885578", "#FAD09F", "#FF8A9A", "#D157A0", "#BEC459", "#456648", "#0086ED", "#886F4C",

              "#34362D", "#B4A8BD", "#00A6AA", "#452C2C", "#636375", "#A3C8C9", "#FF913F", "#938A81",
              "#575329", "#00FECF", "#B05B6F", "#8CD0FF", "#3B9700", "#04F757", "#C8A1A1", "#1E6E00",
              "#7900D7", "#A77500", "#6367A9", "#A05837", "#6B002C", "#772600", "#D790FF", "#9B9700",
              "#549E79", "#FFF69F", "#201625", "#72418F", "#BC23FF", "#99ADC0", "#3A2465", "#922329",
              "#5B4534", "#FDE8DC", "#404E55", "#0089A3", "#CB7E98", "#A4E804", "#324E72", "#6A3A4C",
              "#83AB58", "#001C1E", "#D1F7CE", "#004B28", "#C8D0F6", "#A3A489", "#806C66", "#222800",
              "#BF5650", "#E83000", "#66796D", "#DA007C", "#FF1A59", "#8ADBB4", "#1E0200", "#5B4E51",
              "#C895C5", "#320033", "#FF6832", "#66E1D3", "#CFCDAC", "#D0AC94", "#7ED379", "#012C58"
              ]

    _, axs = plt.subplots(len(gammas), 1)
    ax_ind = -1
    for gamma, y_alig in zip(gammas, ys_alig):
        ax_ind += 1
        label_counter = -1

        for i in range(gamma.shape[1]):
            current_color = colors[i]
            axs[ax_ind].plot(gamma[:, i], color=current_color, linewidth=0.5)

            if y_alig is not None:
                if zero_ones_labels:
                    # If we use zero_ones labels, dashed lines may overlap which makes discriminating
                    # overlapped ground truth labels impossible.
                    if torch.sum(y_alig[:, i] == 1.0) > 0:
                        label_counter += 1

                    axs[ax_ind].plot(torch.clamp_min_(y_alig[:, i] - 0.04 * label_counter, 0.0), '--',
                                     color=current_color, linewidth=0.5)
                else:
                    axs[ax_ind].plot(y_alig[:, i], '--', color=current_color, linewidth=0.5)
                for j in range(gamma.shape[0]):
                    # This line makes the overlapped colors in case of ZERO_ONES, and white in case of
                    # prob. labels.
                    if y_alig[j][i] == 1:
                        axs[ax_ind].axvspan(j - 0.5, j + 0.5, facecolor=current_color, alpha=0.2)


def plot_files(dset, indices, processed=False, zero_ones_labels=False, plot_qinit_only=False):
    """
    If processed is true, it plots the exactly same gammas that are on input to the loss fn.
    """
    with torch.no_grad():
        plt.clf()
        gammas = []
        ys_alig = []
        for sp in indices:
            _, _, fn, x, qinit, y = dset[sp]

            if plot_qinit_only:
                gamma = qinit
            else:
                pi_init = torch.ones(qinit.size()[1], dtype=torch.double, requires_grad=False) / qinit.size()[1]

                if args.run_dist:
                    log_gamma, gamma, pi, Li = ddp_model(x, pi=pi_init, gamma=qinit)
                    if processed:
                        gamma = ddp_model.module.process_gammas(log_gamma)
                else:
                    log_gamma, gamma, pi, Li = model(x, pi=pi_init, gamma=qinit)
                    if processed:
                        gamma = model.process_gammas(log_gamma)

            decisions_t = torch.transpose(gamma, 0, 1)
            cost_mx2 = torch.matmul(1 - decisions_t, y) + torch.matmul(decisions_t, 1 - y)
            pred_alig, ref_alig = linear_sum_assignment(cost_mx2.detach())
            y_alig = y[:, ref_alig]

            gammas.append(gamma)
            ys_alig.append(y_alig)

        plot_gammas(gammas, ys_alig, zero_ones_labels)


def eval_model(model, dset, run_dist, der_wrapper: DER, rank, world_size):
    model.eval()
    with torch.no_grad():
        all_ders = []

        for batch in dset:
            batch_der = []
            for sd, sn, fn, x, qinit, y in batch:
                pi_init = torch.ones(qinit.size()[1],
                                     dtype=torch.double,
                                     requires_grad=False) / qinit.size()[1]

                log_gamma, gamma, pi, Li = model(x, pi=pi_init, gamma=qinit)

                assert (np.all(sd[fn][0] == np.array(sn)))
                start, end = sd[fn][1].T

                labels1st = torch.argmax(gamma, dim=1).numpy()
                der, speech_amount = der_wrapper.der_tensor(start, end, labels1st, fn)
                batch_der.append((fn, der, speech_amount))

            if run_dist:
                batch_ders = [[] for _ in range(world_size)]
                dist.all_gather_object(batch_ders, batch_der)
            else:
                batch_ders = [batch_der]

            if rank == 0:
                all_ders.append(batch_ders)

        if rank == 0:
            # Get rid of repetitive files (sampling w repetitions is used for dist. batching)
            files_dict = dict()
            for batch_ders in all_ders:
                for batch_der in batch_ders:
                    for fn, der, speech_amount in batch_der:
                        files_dict[fn] = (der * speech_amount, speech_amount)

            total_der = sum([files_dict[x][0] for x in files_dict])
            total_speech = sum([files_dict[x][1] for x in files_dict])
            return 100.0 * total_der / total_speech


def get_exp_name(args):
    suffix = datetime.strftime(datetime.now(), '%b%d_%H-%M-%S_') + gethostname()

    if 'dihard' in args.in_trainlist.lower():
        if 'domains' in args.in_trainlist.lower():
            train_dset = 'dihard_' + '.'.join(args.in_trainlist.lower().split('/')[1].split('.')[:-2])
        else:
            train_dset = 'dihard'
    elif 'ch' in args.in_trainlist.lower() or '/train.txt' in args.in_trainlist.lower():
        train_dset = 'callhome'
    elif 'ami' in args.in_trainlist.lower():
        train_dset = 'ami'
    else:
        train_dset = os.path.basename(args.in_trainlist)

    config_name = os.path.basename(args.config_path[0].abs_path).split('.yml')[0]
    back_prop_type_initials = ''.join(map(lambda x: x[0], args.backprop_type.name.lower().split('_')))

    name = f"{config_name}" \
           f"__{train_dset}__" \
           f"{args.gt_label_type}" \
           f"_fa_{args.Fa}" \
           f"_fb_{args.Fb}" \
           f"_pl_{args.loopP}" \
           f"_loss_{args.loss.name.lower()}" \
           f"_ils_{args.initial_loss_scale}" \
           f"_mit_{args.max_iters}" \
           f"_bpt_{back_prop_type_initials}" \
           f"_brc_{args.regularization_coeff_eb}" \
           f"_wrc_{args.regularization_coeff_ew}" \
           f"_ls_{int(args.use_loss_scale)}" \
           f"_reg_{int(args.use_regularization)}" \
           f"_ism_{float(args.init_smoothing)}"
    if args.use_gmm:
        name += '_gmm'
    else:
        name += '_hmm'

    if args.exp_name_tag is not None:
        name += '_tag_' + args.exp_name_tag

    return f'{name}_{suffix}'


if __name__ == '__main__':
    parser = yamlargparse.ArgumentParser()
    parser.add_argument('--in-GTlabels-dir', required=True, type=str,
                        help='Directory with the ground truth xvector labels')
    parser.add_argument('--in-INITlabels-dir', required=True, type=str,
                        help='Directory with the AHC initialization xvector labels')
    parser.add_argument('--in-val-GTlabels-dir', required=False, type=str,
                        help='Directory with the ground truth xvector labels for validation set')
    parser.add_argument('--in-val-INITlabels-dir', required=False, type=str,
                        help='Directory with the AHC initialization xvector labels for validation set')
    parser.add_argument('--in-trainlist', required=True, type=str,
                        help='txt list of training files')
    parser.add_argument('--in-vallist', required=True, type=str,
                        help='txt list of validation files')
    parser.add_argument('--in-train-gt-rttm-dir', required=True, type=str,
                        help='Path to the directory containing training ground truth RTTM files.')
    parser.add_argument('--in-val-gt-rttm-dir', required=True, type=str,
                        help='Path to the directory containing validation ground truth RTTM files.')
    parser.add_argument('--out-rttm-dir', required=True, type=str,
                        help='Directory to store output rttm files')
    parser.add_argument('--xvec-ark-dir', required=True, type=str,
                        help='Kaldi ark with x-vectors from one or more '
                             'input recordings. Attention: all x-vectors from '
                             'one recording must be in one ark file')
    parser.add_argument('--segments-dir', required=True, type=str,
                        help='File with x-vector timing info. See '
                             'diarization_lib.read_xvector_timing_dict')
    parser.add_argument('--val-xvec-ark-dir', required=False, type=str,
                        help='Kaldi ark with x-vectors from one or more '
                             'input recordings. Attention: all x-vectors from '
                             'one recording must be in one ark file')
    parser.add_argument('--val-segments-dir', required=False, type=str,
                        help='File with x-vector timing info. See '
                             'diarization_lib.read_xvector_timing_dict')
    # Eval Dataset Params
    parser.add_argument('--in-eval-list', required=False, type=str,
                        help='txt list of evaluation files')
    parser.add_argument('--in-eval-GTlabels-dir', required=False, type=str,
                        help='Directory with the ground truth xvector labels')
    parser.add_argument('--in-eval-INITlabels-dir', required=False, type=str,
                        help='Directory with the AHC initialization xvector labels for evaluation set.')
    parser.add_argument('--eval-xvec-ark-dir', required=False, type=str,
                        help='Kaldi ark with x-vectors from one or more '
                             'input recordings. Attention: all x-vectors from '
                             'one recording must be in one ark file')
    parser.add_argument('--eval-segments-dir', required=False, type=str,
                        help='File with x-vector timing info. See '
                             'diarization_lib.read_xvector_timing_dict')
    parser.add_argument('--in-eval-gt-rttm-dir', required=False, type=str,
                        help='Path to the directory containing evaluation ground truth RTTM files.')

    parser.add_argument('--xvec-transform', required=True, type=str,
                        help='path to x-vector transformation h5 file')
    parser.add_argument('--plda-file', required=True, type=str,
                        help='File with PLDA model in Kaldi format used for '
                             'AHC and VB-HMM x-vector clustering')
    parser.add_argument('--threshold', required=True, type=float,
                        help='args.threshold (bias) used for AHC')
    parser.add_argument('--lda-dim', required=True, type=int,
                        help='For VB-HMM, x-vectors are reduced to this '
                             'dimensionality using LDA')
    parser.add_argument('--Fa', required=True, type=float,
                        help='Parameter of VB-HMM (see VBx.VBx)')
    parser.add_argument('--Fb', required=True, type=float,
                        help='Parameter of VB-HMM (see VBx.VBx)')
    parser.add_argument('--loopP', required=True, type=float,
                        help='Parameter of VB-HMM (see VBx.VBx)')
    parser.add_argument('--target-energy', required=False, type=float,
                        default=1.0,
                        help='Parameter affecting AHC if the similarity '
                             'matrix is obtained with PLDA. See '
                             'diarization_lib.kaldi_ivector_plda_scoring_dense')
    parser.add_argument('--init-smoothing', required=False, type=float,
                        default=3,
                        help='AHC produces hard assignments of x-vetors to '
                             'speakers. These are "smoothed" to soft '
                             'assignments as the args.initialization for '
                             'VB-HMM. This parameter controls the amount of '
                             'smoothing. Not so important, high value '
                             '(e.g. 10) is OK  => keeping hard assigment')
    parser.add_argument('--gt-label-type', required=False, type=str,
                        default="SINGLE_SPK",
                        help='Ground truth labels are used as probabilities '
                             'that denote the percentage of presence of each '
                             'speaker\'s speech in the x-vector (PROBABILITIES), '
                             'as 0 or 1 values that denote the presence or absence '
                             'of speech of a speaker in a x-vector (ZERO_ONES) '
                             'or as a indicator of the speaker that speaks the most '
                             'in the x-vector (SINGLE_SPK)')
    parser.add_argument('-c', '--config-path', required=True, type=str, action=yamlargparse.ActionConfigFile,
                        help='Path to config.json file.')
    parser.add_argument('--run-dist', required=False, default=False, action='store_true',
                        help='If present, the training has to be run with torchrun.')
    parser.add_argument('--continue-log-dir', required=False,
                        help='Tensorboard log dir to continue the training process from.')
    parser.add_argument('--eval-after-steps', required=False, type=int, default=-1,
                        help='Number of update steps (batches) to evaluate the model on validation data after.')
    parser.add_argument('--eval-after-epochs', required=False, type=int, default=20,
                        help='Number of epochs after which the model is going to be evaluated on train, dev, eval'
                             'datasets. This argument can be overridden by --eval-after-steps, which means that '
                             'the model will be evaluated after --eval-after-epochs number of epochs only if '
                             '--eval-after-steps is not present (i.e. set to -1).')
    parser.add_argument('--save-checkpoint-after-steps', required=False, type=int, default=-1,
                        help='Saves model checkpoint after n update steps (batches).')
    parser.add_argument('--tb-path', required=False, type=str,
                        help='Path to the tensorboard logging dir. If not provided, a default one will be used, '
                             'which consists of date and hostname.')
    parser.add_argument('--num-threads-per-worker', type=int, default=1,
                        help='Number of torch threads used for tensor operations.')
    parser.add_argument('--use-gmm', required=False, default=False, action='store_true',
                        help='If present, GMM is used instead of HMM, i.e. ploop=0.')
    parser.add_argument('--plot-gammas', required=False, default=False, action='store_true',
                        help='Gamma plots of pre-selected files are logged into TensorBoard if the option is present.')
    parser.add_argument('--exp-name-tag', required=False, type=str,
                        help='It helps with classifying the experiments into groups/tags to better navigate through '
                             'TB search.')
    parser.add_argument('--init-model-path', required=False, type=str,
                        help='Path to initial model. If loaded, Fa, Fb, loopP will be overridden.')

    # YAML config parameters
    parser.add_argument('--avg-last-n-iters', type=int, default=-1,
                        help='Dictates how many last VB iterations are averaged to compute gradients '
                             '(in case of using after_each_iter backprop_type)')
    parser.add_argument('--backprop-type', type=str, default='after_each_iter',
                        help='Gradient computation type: after_convergence, or after_each_iter.')
    parser.add_argument('--batch-size', type=int, default=8, help='Batch size.')
    parser.add_argument('--early-stop-vb', type=bool, default=False,
                        help='If True, VB inference is stopped during training if ELBO stops increasing.')
    parser.add_argument('--eval-max-iters', type=int, default=40,
                        help='Max number of VB iterations used during inference/evaluation.')
    parser.add_argument('--eval-early-stop-vb', type=bool, default=True,
                        help='Same as --early-stop-vb but during evaluation.')
    parser.add_argument('--epochs', type=int, default=500, help='Number of passes through the whole training set.')
    parser.add_argument('--initial-loss-scale', type=float, default=1,
                        help='Initial value of loss scale (calibration).')
    parser.add_argument('--loss', type=str, default='EDE', help='Training loss, either BCE or DER.')
    parser.add_argument('--lr', help='Learning rates.')
    parser.add_argument('--max-iters', type=int, default=10,
                        help='Max number of VB iterations during training (if early_stop_vb is False, than it\'')
    parser.add_argument('--regularization-coeff-eb', default=0, type=float,
                        help='KL divergence between-speaker covariance matrix regularization coefficient.')
    parser.add_argument('--regularization-coeff-ew', default=0, type=float,
                        help='KL divergence within-speaker covariance matrix regularization coefficient.')
    parser.add_argument('--trainable-params', default=['fa', 'fb', 'lp', 'ism'], nargs='*',
                        help='List of trainable parameters.')
    parser.add_argument('--use-full-tr', type=bool, default=False,
                        help='Use full transition matrix (i.e. allow HMM to have connection '
                             'between all pairs of nodes).')
    parser.add_argument('--use-loss-scale', type=bool, default=False,
                        help='If true, loss scale (calibration) is used.')
    parser.add_argument('--use-loss-weights', type=bool, default=False,
                        help='Gradients are computed as a weighted average of the VB iteration gradients.')
    parser.add_argument('--use-regularization', type=bool, default=False,
                        help='If True, KL divergence regularization towards the original (generatively trained) '
                             'PLDA covariance matrices is used.')
    parser.add_argument('--use-sigmoid-loss-weights', type=bool, default=False,
                        help='If true, loss weights for iteration i are computed as 1 - product of '
                             'weights for weights i,i+1, ..., n.')

    args = parser.parse_args()

    if args.run_dist:
        dist.init_process_group("gloo")
        rank = dist.get_rank()
        world_size = dist.get_world_size()
        print(f"Start running basic DDP example on rank {rank}.")
    else:
        rank = 0
        world_size = 1

    # Continue training, i.e. use the previously set config file.
    if args.continue_log_dir is not None:
        config_fname = list(filter(lambda x: '.yml' in x, os.listdir(args.continue_log_dir)))[0]
        args = parser.parse_args(['-c', f'{args.continue_log_dir}/{config_fname}'])

    process_args(args)

    assert 0 <= args.loopP <= 1, f'Expecting loopP between 0 and 1, got {args.loopP} instead.'
    plda_mu, plda_psi, plda_tr = read_kaldi_plda(args.plda_file)

    torch.set_num_threads(args.num_threads_per_worker)

    if rank == 0:
        if args.continue_log_dir is None:
            tb_path = args.tb_path
            if not tb_path:
                tb_path = os.getcwd() + '/runs/' + get_exp_name(args)
            tb_writer = SummaryWriter(log_dir=tb_path)
        else:
            tb_writer = SummaryWriter(log_dir=args.continue_log_dir)

        if args.continue_log_dir is None:
            tb_writer.add_hparams({
                'fa': args.Fa,
                'fb': args.Fb,
                'ploop': args.loopP,
                'max_iters': args.max_iters,
                'loss': args.loss.name,
                'use_full_tr': args.use_full_tr,
                'use_loss_scale': args.use_loss_scale,
                'use_loss_weights': args.use_loss_weights,
                'batch_size': args.batch_size,
                'backprop_type': args.backprop_type.name,
                'init_smoothing': args.init_smoothing,
                'gt_label_type': args.gt_label_type,
                'in_trainlist': args.in_trainlist,
                'use_gmm': args.use_gmm
            }, {
                'der': 0
            })
            shell_copy(args.config_path[0].abs_path, tb_writer.log_dir)

    # We want equal batches across workers.
    assert args.batch_size % world_size == 0, 'Batch size have to be divisible by the number of processes.'

    model = VBx.DiscriminativeVB(device=torch.device("cpu"),
                                 phi=torch.from_numpy(plda_psi[:args.lda_dim].copy()),
                                 tr=torch.from_numpy(plda_tr.copy()),
                                 fa=args.Fa,
                                 fb=args.Fb,
                                 loop_p=args.loopP,
                                 config=args,
                                 init_smoothing=args.init_smoothing,
                                 epsilon=1e-6,
                                 use_gmm=args.use_gmm)

    if args.init_model_path is not None:
        model.load_model(args.init_model_path)

    step_counter = 0
    starting_epoch = 1
    if args.continue_log_dir is not None:
        max_ep = max(*[int(x.split('ep_')[1][:-3]) for x in os.listdir(f'{args.continue_log_dir}/checkpoints')])
        starting_epoch, step_counter = model.load_model(f'{args.continue_log_dir}/checkpoints/model_ep_{max_ep}.pt')
        starting_epoch += 1
        step_counter += 1

    if args.run_dist:
        ddp_model = DDP(model, find_unused_parameters=False)

    train_list = np.loadtxt(args.in_trainlist, dtype=object)
    val_list = np.loadtxt(args.in_vallist, dtype=object)

    der_train = DER(train_list, args.in_train_gt_rttm_dir)
    der_val = DER(val_list, args.in_val_gt_rttm_dir)

    if args.in_eval_gt_rttm_dir:
        eval_list = np.loadtxt(args.in_eval_list, dtype=object)
        der_eval = DER(eval_list, args.in_eval_gt_rttm_dir)
    else:
        eval_list = []
        der_eval = None

    dfl = DataFileLoader(plda_mu,
                         args.segments_dir,
                         args.in_INITlabels_dir,
                         args.xvec_ark_dir,
                         args.in_GTlabels_dir,
                         args.xvec_transform,
                         args.init_smoothing,
                         args.gt_label_type)

    # In case some validation file is not provided, use the val part of the training set.
    if args.in_val_INITlabels_dir is None or \
            args.in_val_GTlabels_dir is None or \
            args.val_xvec_ark_dir is None or \
            args.val_segments_dir is None:
        val_dfl = dfl
    else:
        val_dfl = DataFileLoader(plda_mu,
                                 args.val_segments_dir,
                                 args.in_val_INITlabels_dir,
                                 args.val_xvec_ark_dir,
                                 args.in_val_GTlabels_dir,
                                 args.xvec_transform,
                                 args.init_smoothing,
                                 args.gt_label_type)

    # Eval rttm dir arg presence determines whether to score the Eval dataset or not.
    if args.in_eval_gt_rttm_dir:
        eval_dfl = DataFileLoader(plda_mu,
                                  args.eval_segments_dir,
                                  args.in_eval_INITlabels_dir,
                                  args.eval_xvec_ark_dir,
                                  args.in_eval_GTlabels_dir,
                                  args.xvec_transform,
                                  args.init_smoothing,
                                  args.gt_label_type,
                                  eval_mode=True)
        eval_dset = BasicDataset(eval_list, eval_dfl)
    else:
        eval_dset = None

    train_dset = BasicDataset(train_list, dfl)
    val_dset = BasicDataset(val_list, val_dfl)

    val_steps = args.eval_after_steps
    if val_steps == -1:
        val_steps = ceil(len(train_list) / args.batch_size) * args.eval_after_epochs

    save_steps = args.save_checkpoint_after_steps
    if save_steps == -1:
        save_steps = ceil(len(train_list) / args.batch_size)

    if args.run_dist:
        train_dsampler = DistributedSampler(train_dset,
                                            num_replicas=world_size,
                                            rank=rank,
                                            shuffle=True)
        train_dl = DataLoader(train_dset,
                              batch_size=args.batch_size // world_size,
                              collate_fn=lambda x: x,
                              sampler=train_dsampler)
        val_dl = DataLoader(val_dset,
                            batch_size=args.batch_size // world_size,
                            shuffle=False,
                            collate_fn=lambda x: x,
                            sampler=DistributedSampler(val_dset,
                                                       num_replicas=world_size,
                                                       rank=rank,
                                                       shuffle=False))

        if args.in_eval_gt_rttm_dir:
            eval_dl = DataLoader(eval_dset,
                                 batch_size=args.batch_size // world_size,
                                 shuffle=False,
                                 collate_fn=lambda x: x,
                                 sampler=DistributedSampler(eval_dset,
                                                            num_replicas=world_size,
                                                            rank=rank,
                                                            shuffle=False))
    else:
        train_dl = create_data_loader(train_dset,
                                      batch_size=args.batch_size,
                                      shuffle=True)
        val_dl = create_data_loader(val_dset,
                                    batch_size=args.batch_size,
                                    shuffle=False)

        if args.in_eval_gt_rttm_dir:
            eval_dl = create_data_loader(eval_dset,
                                         batch_size=args.batch_size,
                                         shuffle=False)

    if args.run_dist:
        train_der = eval_model(ddp_model, train_dl, args.run_dist, der_train, rank, world_size)
        val_der = eval_model(ddp_model, val_dl, args.run_dist, der_val, rank, world_size)

        if args.in_eval_gt_rttm_dir:
            eval_der = eval_model(ddp_model, eval_dl, args.run_dist, der_eval, rank, world_size)
    else:
        train_der = eval_model(model, train_dl, args.run_dist, der_train, rank, world_size)
        val_der = eval_model(model, val_dl, args.run_dist, der_val, rank, world_size)

        if args.in_eval_gt_rttm_dir:
            eval_der = eval_model(model, eval_dl, args.run_dist, der_eval, rank, world_size)

    # Plot init gammas
    if rank == 0:
        if not args.continue_log_dir:
            tb_writer.add_scalar('der/train_ep', train_der, 0)
            tb_writer.add_scalar('der/val_ep', val_der, 0)

            if args.in_eval_gt_rttm_dir:
                tb_writer.add_scalar('der/eval_ep', eval_der, 0)

        if 'dihard_exp' in args.in_GTlabels_dir:
            train_file_indices = [78, 35, 90]
            val_file_indices = [2, 5, 1, 3, 4]
        elif 'ch' in args.in_GTlabels_dir:
            train_file_indices = [1, 5, 185]
            val_file_indices = [2, 5, 1, 3, 4]
        elif 'ami' in args.in_GTlabels_dir:
            train_file_indices = [1, 2, 3]
            val_file_indices = [2, 5, 1, 3, 4]
        else:
            train_file_indices = [0, 1, 2, 3]
            val_file_indices = [0, 1, 2, 3]

        zero_ones_labels = args.gt_label_type == 'ZERO_ONES'

        if args.plot_gammas:
            plot_files(train_dset, train_file_indices, plot_qinit_only=True, zero_ones_labels=zero_ones_labels)
            tb_writer.add_figure(f'gammas/train_init_gammas', plt.gcf(), 0)
            plot_files(train_dset, val_file_indices, plot_qinit_only=True, zero_ones_labels=zero_ones_labels)
            tb_writer.add_figure(f'gammas/val_init_gammas', plt.gcf(), 0)

        # Save starting point model.
        save_path = tb_writer.log_dir
        if not os.path.exists(f'{save_path}/checkpoints'):
            os.mkdir(f'{save_path}/checkpoints')

        if args.run_dist:
            ddp_model.module.save_model(f'{save_path}/checkpoints/model_ep_0.pt', 0, 0)
        else:
            model.save_model(f'{save_path}/checkpoints/model_ep_0.pt', 0, 0)

    if args.run_dist:
        dist.barrier()

    for ep in range(starting_epoch, starting_epoch + args.epochs):
        if args.run_dist:
            train_dsampler.set_epoch(ep)

        ep_losses = []
        ep_raw_losses = []
        for batch in train_dl:
            step_counter += 1

            batch_losses = []
            batch_raw_losses = []
            for _, _, _, x, qinit, y in batch:
                pi_init = torch.ones(qinit.size()[1],
                                     dtype=torch.double,
                                     requires_grad=False) / qinit.size()[1]

                if args.run_dist:
                    ddp_model.module.train()
                    log_gamma, gamma, pi, Li = ddp_model(x, pi=pi_init, gamma=qinit)
                else:
                    model.train()
                    log_gamma, gamma, pi, Li = model(x, pi=pi_init, gamma=qinit)

                loss, raw_loss = model.compute_loss(y)
                batch_losses.append(loss)
                batch_raw_losses.append(raw_loss)
                ep_losses.append(loss.detach())
                ep_raw_losses.append(raw_loss)

            avg_loss = sum(batch_losses) / len(batch_losses)
            avg_raw_loss = sum(batch_raw_losses) / len(batch_raw_losses)
            model.optimize_params(avg_loss)

            if rank == 0:
                model.tb_log_params(tb_writer, 'params', 'step', step_counter)
                model.tb_log_params_grads(tb_writer, 'grads', 'step', step_counter)

                if args.run_dist:
                    KL_Ew, KL_Eb = ddp_model.module.compute_KL()
                else:
                    KL_Ew, KL_Eb = model.compute_KL()

                tb_writer.add_scalar('params/KL_Eb_step', KL_Eb.item(), step_counter)
                tb_writer.add_scalar('params/KL_Ew_step', KL_Ew.item(), step_counter)

            if args.run_dist:
                batch_avg_losses = [0 for _ in range(world_size)]
                batch_avg_raw_losses = [0 for _ in range(world_size)]
                dist.all_gather_object(batch_avg_losses, avg_loss.item())
                dist.all_gather_object(batch_avg_raw_losses, avg_raw_loss.item())
            else:
                batch_avg_losses = [avg_loss]
                batch_avg_raw_losses = [avg_raw_loss]

            if rank == 0:
                batch_avg_loss = sum(batch_avg_losses) / len(batch_avg_losses)
                batch_avg_raw_loss = sum(batch_avg_raw_losses) / len(batch_avg_raw_losses)
                print(f"BatchLoss: {sum(batch_avg_losses) / len(batch_avg_losses)}")
                tb_writer.add_scalar('train/loss_step', batch_avg_loss, step_counter)
                tb_writer.add_scalar('train/raw_loss_step', batch_avg_raw_loss, step_counter)

            if step_counter % val_steps == 0:
                if args.run_dist:
                    ddp_model.module.eval()
                else:
                    model.eval()

                val_ep = step_counter // val_steps
                val_losses = []
                val_raw_losses = []
                with torch.no_grad():
                    for batch in val_dl:
                        for _, _, _, x, qinit, y in batch:
                            pi_init = torch.ones(qinit.size()[1],
                                                 dtype=torch.double,
                                                 requires_grad=False) / qinit.size()[1]

                            if args.run_dist:
                                log_gamma, gamma, pi, Li = ddp_model(x, pi=pi_init, gamma=qinit)
                            else:
                                log_gamma, gamma, pi, Li = model(x, pi=pi_init, gamma=qinit)

                            loss, raw_loss = model.compute_loss(y)
                            val_losses.append(loss)
                            val_raw_losses.append(raw_loss)

                if args.run_dist:
                    dist.barrier()

                if args.run_dist:
                    val_avg_losses = [None for _ in range(world_size)]
                    val_avg_raw_losses = [None for _ in range(world_size)]
                    dist.all_gather_object(val_avg_losses, torch.mean(torch.stack(val_losses)))
                    dist.all_gather_object(val_avg_raw_losses, torch.mean(torch.stack(val_raw_losses)))
                else:
                    val_avg_losses = [torch.mean(torch.stack(val_losses))]
                    val_avg_raw_losses = [torch.mean(torch.stack(val_raw_losses))]

                if args.run_dist:
                    train_der = eval_model(ddp_model, train_dl, args.run_dist, der_train, rank, world_size)
                    val_der = eval_model(ddp_model, val_dl, args.run_dist, der_val, rank, world_size)

                    if args.in_eval_gt_rttm_dir:
                        eval_der = eval_model(ddp_model, eval_dl, args.run_dist, der_eval, rank, world_size)
                else:
                    train_der = eval_model(model, train_dl, args.run_dist, der_train, rank, world_size)
                    val_der = eval_model(model, val_dl, args.run_dist, der_val, rank, world_size)

                    if args.in_eval_gt_rttm_dir:
                        eval_der = eval_model(model, eval_dl, args.run_dist, der_eval, rank, world_size)

                if rank == 0:
                    # ToDo: Mention this val_ep conting into README !!! It can confuse a new user once eval after ..
                    #   flag is set.
                    tb_writer.add_scalar('der/train_ep', train_der, val_ep)
                    tb_writer.add_scalar('der/val_ep', val_der, val_ep)

                    if args.in_eval_gt_rttm_dir:
                        tb_writer.add_scalar('der/eval_ep', eval_der, val_ep)

                    if args.plot_gammas:
                        plot_files(train_dset, train_file_indices, zero_ones_labels=zero_ones_labels)
                        tb_writer.add_figure(f'gammas/train_ep', plt.gcf(), val_ep)
                        plot_files(train_dset, train_file_indices, processed=True, zero_ones_labels=zero_ones_labels)
                        tb_writer.add_figure(f'gammas/train_loss_inp_ep', plt.gcf(), val_ep)

                        plot_files(val_dset, val_file_indices, zero_ones_labels=zero_ones_labels)
                        tb_writer.add_figure(f'gammas/val_ep', plt.gcf(), val_ep)
                        plot_files(val_dset, val_file_indices, processed=True, zero_ones_labels=zero_ones_labels)
                        tb_writer.add_figure(f'gammas/val_loss_inp_ep', plt.gcf(), val_ep)

                    avg_val_loss = sum(val_avg_losses) / len(val_avg_losses)
                    avg_val_raw_loss = sum(val_avg_raw_losses) / len(val_avg_raw_losses)

                    print(f"AvgValLoss: {avg_val_loss}")

                    tb_writer.add_scalar('train/val_loss_epoch', avg_val_loss, val_ep)
                    tb_writer.add_scalar('train/val_raw_loss_epoch', avg_val_raw_loss, val_ep)

            if rank == 0 and step_counter % save_steps == 0:
                save_ep = step_counter // save_steps
                save_path = tb_writer.log_dir
                if not os.path.exists(f'{save_path}/checkpoints'):
                    os.mkdir(f'{save_path}/checkpoints')

                if args.run_dist:
                    ddp_model.module.save_model(f'{save_path}/checkpoints/model_ep_{save_ep}.pt', save_ep,
                                                step_counter)
                else:
                    model.save_model(f'{save_path}/checkpoints/model_ep_{save_ep}.pt', save_ep, step_counter)

        # Log training progress
        if args.run_dist:
            ep_avg_losses = [None for _ in range(world_size)]
            ep_avg_raw_losses = [None for _ in range(world_size)]
            dist.all_gather_object(ep_avg_losses, torch.mean(torch.stack(ep_losses)))
            dist.all_gather_object(ep_avg_raw_losses, torch.mean(torch.stack(ep_raw_losses)))
        else:
            ep_avg_losses = [torch.mean(torch.stack(ep_losses))]
            ep_avg_raw_losses = [torch.mean(torch.stack(ep_raw_losses))]

        if rank == 0:
            avg_ep_loss = sum(ep_avg_losses) / len(ep_avg_losses)
            avg_ep_raw_loss = sum(ep_avg_raw_losses) / len(ep_avg_raw_losses)
            print(f'AvgEpLoss: {avg_ep_loss}')

            tb_writer.add_scalar('train/loss_epoch', avg_ep_loss, ep)
            tb_writer.add_scalar('train/raw_loss_epoch', avg_ep_raw_loss, ep)

        if args.run_dist:
            dist.barrier()
    print("DONE")

    if rank == 0:
        tb_writer.close()

    if args.run_dist:
        dist.destroy_process_group()

vbhmm_train.py
BUTSpeechFIT/DVBx