The official Pytorch implementation of "Frame-wise streaming end-to-end speaker diarization with non-autoregressive self-attention-based attractors".
#!/usr/bin/env python3

# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita)
# Modified by: Yexin Yang
# Modified by: Di Liang
# Licensed under the MIT license.
#
import torch
import numpy as np
from .kaldi_data import KaldiData
from .feature import *

def _count_frames(data_len, size, step):
    # no padding at edges, last remaining samples are ignored
    return int((data_len - size + step) / step)


def _gen_frame_indices(
        data_length, size=2000, step=2000,
        use_last_samples=False,
        label_delay=0,
        subsampling=1):
    i = -1
    for i in range(_count_frames(data_length, size, step)):
        yield i * step, i * step + size
    if use_last_samples and i * step + size < data_length:
        if data_length - (i + 1) * step - label_delay > 0:
            yield (i + 1) * step, data_length


def my_collate(batch):
    data, target, rec = list(zip(*batch))
    return [data, target, rec]


class KaldiDiarizationDataset(torch.utils.data.Dataset):
    def __init__(
            self,
            data_dir,
            chunk_size=2000,
            context_size=0,
            frame_size=1024,
            frame_shift=256,
            subsampling=1,
            rate=16000,
            input_transform=None,
            use_last_samples=False,
            label_delay=0,
            n_speakers=None,
            shuffle=False
            ):
        self.data_dir = data_dir
        self.chunk_size = chunk_size
        self.context_size = context_size
        self.frame_size = frame_size
        self.frame_shift = frame_shift
        self.subsampling = subsampling
        self.input_transform = input_transform
        self.n_speakers = n_speakers
        self.chunk_indices = []
        self.label_delay = label_delay
        self.shuffle = shuffle

        self.data = KaldiData(self.data_dir)

        # make chunk indices: filepath, start_frame, end_frame
        for rec in self.data.wavs:
            data_len = int(self.data.reco2dur[rec] * rate / frame_shift)
            data_len = int(data_len / self.subsampling)
            for st, ed in _gen_frame_indices(
                    data_len, chunk_size, chunk_size, use_last_samples,
                    label_delay=self.label_delay,
                    subsampling=self.subsampling):
                self.chunk_indices.append(
                        (rec, st * self.subsampling, ed * self.subsampling))
        print(len(self.chunk_indices), " chunks")

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

    def __getitem__(self, i):
        rec, st, ed = self.chunk_indices[i]
        Y, T = get_labeledSTFT(
            self.data,
            rec,
            st,
            ed,
            self.frame_size,
            self.frame_shift,
            self.n_speakers)
        # Y: (frame, num_ceps)
        Y = transform(Y, self.input_transform)
        # Y_spliced: (frame, num_ceps * (context_size * 2 + 1))
        Y_spliced = splice(Y, self.context_size)
        # Y_ss: (frame / subsampling, num_ceps * (context_size * 2 + 1))
        Y_ss, T_ss = subsample(Y_spliced, T, self.subsampling)

        Y_ss = torch.from_numpy(Y_ss).float()
        T_ss = torch.from_numpy(T_ss).float()
        
        if self.shuffle:
            order = np.arange(Y_ss.shape[0])
            np.random.shuffle(order)
            Y_ss = Y_ss[order]
            T_ss = T_ss[order]
        
        return Y_ss, T_ss, rec
    
    def __getfulllabel__(self, i):
        rec, st, ed = self.chunk_indices[i]
        Y, T = get_labeledSTFT(
            self.data,
            rec,
            st,
            ed,
            self.frame_size,
            self.frame_shift,
            self.n_speakers)
        # Y: (frame, num_ceps)
        Y = transform(Y, self.input_transform)
        # Y_spliced: (frame, num_ceps * (context_size * 2 + 1))
        Y_spliced = splice(Y, self.context_size)
        # Y_ss: (frame / subsampling, num_ceps * (context_size * 2 + 1))
        Y_ss, T_ss = subsample(Y_spliced, T, self.subsampling)
        T_ss = torch.from_numpy(T_ss).float()
        T = torch.from_numpy(T).float()
        return T, rec
    
    def __get_len__(self, i):
        rec, st, ed = self.chunk_indices[i]
        Y, T = get_labeledSTFT(
            self.data,
            rec,
            st,
            ed,
            self.frame_size,
            self.frame_shift,
            self.n_speakers)
        T = torch.from_numpy(T).float()
        return T.shape[0]

datasets/diarization_dataset.py
Audio-WestlakeU/FS-EEND
# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita)
# Licensed under the MIT license.
# Modified by: Di Liang
# This module is for computing audio features

import numpy as np
import librosa
import scipy.signal

def get_input_dim(
        frame_size,
        context_size,
        transform_type,
        ):
    if transform_type.startswith('logmel23'):
        frame_size = 23
    else:
        fft_size = 1 << (frame_size-1).bit_length()
        frame_size = int(fft_size / 2) + 1
    input_dim = (2 * context_size + 1) * frame_size
    return input_dim


def transform(
        Y,
        transform_type=None,
        dtype=np.float32):
    """ Transform STFT feature

    Args:
        Y: STFT
            (n_frames, n_bins)-shaped np.complex array
        transform_type:
            None, "log"
        dtype: output data type
            np.float32 is expected
    Returns:
        Y (numpy.array): transformed feature
    """
    Y = np.abs(Y)
    if not transform_type:
        pass
    elif transform_type == 'log':
        Y = np.log(np.maximum(Y, 1e-10))
    elif transform_type == 'logmel':
        n_fft = 2 * (Y.shape[1] - 1)
        sr = 16000
        n_mels = 40
        mel_basis = librosa.filters.mel(sr, n_fft, n_mels)
        Y = np.dot(Y ** 2, mel_basis.T)
        Y = np.log10(np.maximum(Y, 1e-10))
    elif transform_type == 'logmel23':
        n_fft = 2 * (Y.shape[1] - 1)
        sr = 8000
        n_mels = 23
        mel_basis = librosa.filters.mel(sr, n_fft, n_mels)
        Y = np.dot(Y ** 2, mel_basis.T)
        Y = np.log10(np.maximum(Y, 1e-10))
    elif transform_type == 'logmel23_mn':
        n_fft = 2 * (Y.shape[1] - 1)
        sr = 8000
        n_mels = 23
        mel_basis = librosa.filters.mel(sr, n_fft, n_mels)
        Y = np.dot(Y ** 2, mel_basis.T)
        Y = np.log10(np.maximum(Y, 1e-10))
        mean = np.mean(Y, axis=0)
        Y = Y - mean
    elif transform_type == 'logmel23_swn':
        n_fft = 2 * (Y.shape[1] - 1)
        sr = 8000
        n_mels = 23
        mel_basis = librosa.filters.mel(sr, n_fft, n_mels)
        Y = np.dot(Y ** 2, mel_basis.T)
        Y = np.log10(np.maximum(Y, 1e-10))
        #b = np.ones(300)/300
        #mean = scipy.signal.convolve2d(Y, b[:, None], mode='same')
        #
        # simple 2-means based threshoding for mean calculation
        powers = np.sum(Y, axis=1)
        th = (np.max(powers) + np.min(powers))/2.0
        for i in range(10):
            th = (np.mean(powers[powers >= th]) + np.mean(powers[powers < th])) / 2
        mean = np.mean(Y[powers > th,:], axis=0)
        Y = Y - mean
    elif transform_type == 'logmel23_mvn':
        n_fft = 2 * (Y.shape[1] - 1)
        sr = 8000
        n_mels = 23
        mel_basis = librosa.filters.mel(sr, n_fft, n_mels)
        Y = np.dot(Y ** 2, mel_basis.T)
        Y = np.log10(np.maximum(Y, 1e-10))
        mean = np.mean(Y, axis=0)
        Y = Y - mean
        std = np.maximum(np.std(Y, axis=0), 1e-10)
        Y = Y / std
    else:
        raise ValueError('Unknown transform_type: %s' % transform_type)
    return Y.astype(dtype)


def subsample(Y, T, subsampling=1):
    """ Frame subsampling
    """
    Y_ss = Y[::subsampling]
    T_ss = T[::subsampling]
    return Y_ss, T_ss


def splice(Y, context_size=0):
    """ Frame splicing

    Args:
        Y: feature
            (n_frames, n_featdim)-shaped numpy array
        context_size:
            number of frames concatenated on left-side
            if context_size = 5, 11 frames are concatenated.

    Returns:
        Y_spliced: spliced feature
            (n_frames, n_featdim * (2 * context_size + 1))-shaped
    """
    Y_pad = np.pad(
            Y,
            [(context_size, context_size), (0, 0)],
            'constant')
    Y_spliced = np.lib.stride_tricks.as_strided(
            Y_pad,
            (Y.shape[0], Y.shape[1] * (2 * context_size + 1)),
            (Y.itemsize * Y.shape[1], Y.itemsize), writeable=False)
    return Y_spliced


def stft(
        data,
        frame_size=1024,
        frame_shift=256):
    """ Compute STFT features

    Args:
        data: audio signal
            (n_samples,)-shaped np.float32 array
        frame_size: number of samples in a frame (must be a power of two)
        frame_shift: number of samples between frames

    Returns:
        stft: STFT frames
            (n_frames, n_bins)-shaped np.complex64 array
    """
    # round up to nearest power of 2
    fft_size = 1 << (frame_size-1).bit_length()
    # HACK: The last frame is ommited
    #       as librosa.stft produces such an excessive frame
    if len(data) % frame_shift == 0:
        return librosa.stft(data, n_fft=fft_size, win_length=frame_size,
                            hop_length=frame_shift).T[:-1]
    else:
        return librosa.stft(data, n_fft=fft_size, win_length=frame_size,
                            hop_length=frame_shift).T


def _count_frames(data_len, size, shift):
    # HACK: Assuming librosa.stft(..., center=True)
    n_frames = 1 + int(data_len / shift)
    if data_len % shift == 0:
        n_frames = n_frames - 1
    return n_frames


def get_frame_labels(
        kaldi_obj,
        rec,
        start=0,
        end=None,
        frame_size=1024,
        frame_shift=256,
        n_speakers=None):
    """ Get frame-aligned labels of given recording
    Args:
        kaldi_obj (KaldiData)
        rec (str): recording id
        start (int): start frame index
        end (int): end frame index
            None means the last frame of recording
        frame_size (int): number of frames in a frame
        frame_shift (int): number of shift samples
        n_speakers (int): number of speakers
            if None, the value is given from data
    Returns:
        T: label
            (n_frames, n_speakers)-shaped np.int32 array
    """
    filtered_segments = kaldi_obj.segments[kaldi_obj.segments['rec'] == rec]
    speakers = np.unique(
            [kaldi_obj.utt2spk[seg['utt']] for seg
                in filtered_segments]).tolist()
    if n_speakers is None:
        n_speakers = len(speakers)
    es = end * frame_shift if end is not None else None
    data, rate = kaldi_obj.load_wav(
            rec, start * frame_shift, es)
    n_frames = _count_frames(len(data), frame_size, frame_shift)
    T = np.zeros((n_frames, n_speakers), dtype=np.int32)
    if end is None:
        end = n_frames

    for seg in filtered_segments:
        speaker_index = speakers.index(kaldi_obj.utt2spk[seg['utt']])
        start_frame = np.rint(
                seg['st'] * rate / frame_shift).astype(int)
        end_frame = np.rint(
                seg['et'] * rate / frame_shift).astype(int)
        rel_start = rel_end = None
        if start <= start_frame and start_frame < end:
            rel_start = start_frame - start
        if start < end_frame and end_frame <= end:
            rel_end = end_frame - start
        if rel_start is not None or rel_end is not None:
            T[rel_start:rel_end, speaker_index] = 1
    return T


def get_labeledSTFT(
        kaldi_obj,
        rec, start, end, frame_size, frame_shift,
        n_speakers=None,
        use_speaker_id=False):
    """ Extracts STFT and corresponding labels

    Extracts STFT and corresponding diarization labels for
    given recording id and start/end times

    Args:
        kaldi_obj (KaldiData)
        rec (str): recording id
        start (int): start frame index
        end (int): end frame index
        frame_size (int): number of samples in a frame
        frame_shift (int): number of shift samples
        n_speakers (int): number of speakers
            if None, the value is given from data
    Returns:
        Y: STFT
            (n_frames, n_bins)-shaped np.complex64 array,
        T: label
            (n_frmaes, n_speakers)-shaped np.int32 array.
    """
    data, rate = kaldi_obj.load_wav(
            rec, start * frame_shift, end * frame_shift)
    Y = stft(data, frame_size, frame_shift)
    filtered_segments = kaldi_obj.segments[rec]
    # filtered_segments = kaldi_obj.segments[kaldi_obj.segments['rec'] == rec]
    speakers = np.unique(
            [kaldi_obj.utt2spk[seg['utt']] for seg
                in filtered_segments]).tolist()
    if n_speakers is None:
        n_speakers = len(speakers)
    T = np.zeros((Y.shape[0], n_speakers), dtype=np.int32)

    if use_speaker_id:
        all_speakers = sorted(kaldi_obj.spk2utt.keys())
        S = np.zeros((Y.shape[0], len(all_speakers)), dtype=np.int32)

    for seg in filtered_segments:
        speaker_index = speakers.index(kaldi_obj.utt2spk[seg['utt']])
        if use_speaker_id:
            all_speaker_index = all_speakers.index(kaldi_obj.utt2spk[seg['utt']])
        start_frame = np.rint(
                seg['st'] * rate / frame_shift).astype(int)
        end_frame = np.rint(
                seg['et'] * rate / frame_shift).astype(int)
        rel_start = rel_end = None
        if start <= start_frame and start_frame < end:
            rel_start = start_frame - start
        if start < end_frame and end_frame <= end:
            rel_end = end_frame - start
        if rel_start is not None or rel_end is not None:
            T[rel_start:rel_end, speaker_index] = 1
            if use_speaker_id:
                S[rel_start:rel_end, all_speaker_index] = 1

    if use_speaker_id:
        return Y, T, S
    else:
        return Y, T

def get_ET(
        kaldi_obj,
        rec, start, end, frame_size, frame_shift, subsampling, rate,
        n_speakers=None,
        use_speaker_id=False):
    E = kaldi_obj.load_emb(rec, int(start / subsampling), int(end / subsampling))
    filtered_segments = kaldi_obj.segments[rec]
    # filtered_segments = kaldi_obj.segments[kaldi_obj.segments['rec'] == rec]
    speakers = np.unique(
            [kaldi_obj.utt2spk[seg['utt']] for seg
                in filtered_segments]).tolist()
    if n_speakers is None:
        n_speakers = len(speakers)
    T = np.zeros((E.shape[0] * subsampling, n_speakers), dtype=np.int32)

    if use_speaker_id:
        all_speakers = sorted(kaldi_obj.spk2utt.keys())
        S = np.zeros((E.shape[0] * subsampling, len(all_speakers)), dtype=np.int32)

    for seg in filtered_segments:
        speaker_index = speakers.index(kaldi_obj.utt2spk[seg['utt']])
        if use_speaker_id:
            all_speaker_index = all_speakers.index(kaldi_obj.utt2spk[seg['utt']])
        start_frame = np.rint(
                seg['st'] * rate / frame_shift).astype(int)
        end_frame = np.rint(
                seg['et'] * rate / frame_shift).astype(int)
        rel_start = rel_end = None
        if start <= start_frame and start_frame < end:
            rel_start = start_frame - start
        if start < end_frame and end_frame <= end:
            rel_end = end_frame - start
        if rel_start is not None or rel_end is not None:
            T[rel_start:rel_end, speaker_index] = 1
            if use_speaker_id:
                S[rel_start:rel_end, all_speaker_index] = 1

    if use_speaker_id:
        return E, T, S
    else:
        return E, T

datasets/feature.py
Audio-WestlakeU/FS-EEND
# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita)
# Licensed under the MIT license.
# Modified by: Di Liang
# This library provides utilities for kaldi-style data directory.


from __future__ import print_function
import os
import sys
import numpy as np
import subprocess
import soundfile as sf
import io
from functools import lru_cache


def load_segments(segments_file):
    """ load segments file as array """
    if not os.path.exists(segments_file):
        return None
    return np.loadtxt(
            segments_file,
            dtype=[('utt', 'object'),
                   ('rec', 'object'),
                   ('st', 'f'),
                   ('et', 'f')],
            ndmin=1)


def load_segments_hash(segments_file):
    ret = {}
    if not os.path.exists(segments_file):
        return None
    for line in open(segments_file):
        utt, rec, st, et = line.strip().split()
        ret[utt] = (rec, float(st), float(et))
    return ret


def load_segments_rechash(segments_file):
    ret = {}
    if not os.path.exists(segments_file):
        return None
    for line in open(segments_file):
        utt, rec, st, et = line.strip().split()
        if rec not in ret:
            ret[rec] = []
        ret[rec].append({'utt':utt, 'st':float(st), 'et':float(et)})
    return ret


def load_wav_scp(wav_scp_file):
    """ return dictionary { rec: wav_rxfilename } """
    lines = [line.strip().split(None, 1) for line in open(wav_scp_file)]
    return {x[0]: x[1] for x in lines}

def load_emb_scp(emb_scp_file):
    """ return dictionary { rec: wav_rxfilename } """
    lines = [line.strip().split(None, 1) for line in open(emb_scp_file)]
    return {x[0]: x[1] for x in lines}

@lru_cache(maxsize=1)
def load_wav(wav_rxfilename, start=0, end=None):
    """ This function reads audio file and return data in numpy.float32 array.
        "lru_cache" holds recently loaded audio so that can be called
        many times on the same audio file.
        OPTIMIZE: controls lru_cache size for random access,
        considering memory size
    """
    if wav_rxfilename.endswith('|'):
        # input piped command
        p = subprocess.Popen(wav_rxfilename[:-1], shell=True,
                             stdout=subprocess.PIPE)
        data, samplerate = sf.read(io.BytesIO(p.stdout.read()),
                                   dtype='float32')
        # cannot seek
        data = data[start:end]
    elif wav_rxfilename == '-':
        # stdin
        data, samplerate = sf.read(sys.stdin, dtype='float32')
        # cannot seek
        data = data[start:end]
    else:
        # normal wav file
        data, samplerate = sf.read(wav_rxfilename, start=start, stop=end)
    return data, samplerate

@lru_cache(maxsize=1)
def load_emb(emb_rxfilename, start=0, end=None):
    data = np.load(emb_rxfilename)
    return data[start:end]

def load_utt2spk(utt2spk_file):
    """ returns dictionary { uttid: spkid } """
    lines = [line.strip().split(None, 1) for line in open(utt2spk_file)]
    return {x[0]: x[1] for x in lines}


def load_spk2utt(spk2utt_file):
    """ returns dictionary { spkid: list of uttids } """
    if not os.path.exists(spk2utt_file):
        return None
    lines = [line.strip().split() for line in open(spk2utt_file)]
    return {x[0]: x[1:] for x in lines}


def load_reco2dur(reco2dur_file):
    """ returns dictionary { recid: duration }  """
    if not os.path.exists(reco2dur_file):
        return None
    lines = [line.strip().split(None, 1) for line in open(reco2dur_file)]
    return {x[0]: float(x[1]) for x in lines}


def process_wav(wav_rxfilename, process):
    """ This function returns preprocessed wav_rxfilename
    Args:
        wav_rxfilename: input
        process: command which can be connected via pipe,
                use stdin and stdout
    Returns:
        wav_rxfilename: output piped command
    """
    if wav_rxfilename.endswith('|'):
        # input piped command
        return wav_rxfilename + process + "|"
    else:
        # stdin "-" or normal file
        return "cat {} | {} |".format(wav_rxfilename, process)


def extract_segments(wavs, segments=None):
    """ This function returns generator of segmented audio as
        (utterance id, numpy.float32 array)
        TODO?: sampling rate is not converted.
    """
    if segments is not None:
        # segments should be sorted by rec-id
        for seg in segments:
            wav = wavs[seg['rec']]
            data, samplerate = load_wav(wav)
            st_sample = np.rint(seg['st'] * samplerate).astype(int)
            et_sample = np.rint(seg['et'] * samplerate).astype(int)
            yield seg['utt'], data[st_sample:et_sample]
    else:
        # segments file not found,
        # wav.scp is used as segmented audio list
        for rec in wavs:
            data, samplerate = load_wav(wavs[rec])
            yield rec, data


class KaldiData:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.segments = load_segments_rechash(
                os.path.join(self.data_dir, 'segments'))
        self.utt2spk = load_utt2spk(
                os.path.join(self.data_dir, 'utt2spk'))
        self.wavs = load_wav_scp(
                os.path.join(self.data_dir, 'wav.scp'))
        self.reco2dur = load_reco2dur(
                os.path.join(self.data_dir, 'reco2dur'))
        self.spk2utt = load_spk2utt(
                os.path.join(self.data_dir, 'spk2utt'))
        # self.embs = load_emb_scp(
        #         os.path.join(self.data_dir, 'emb.scp'))

    def load_wav(self, recid, start=0, end=None):
        data, rate = load_wav(
            self.wavs[recid], start, end)
        return data, rate

    def load_emb(self, recid, start=0, end=None):
        emb = load_emb(
            self.embs[recid], start, end)
        return emb
datasets/kaldi_data.py
Audio-WestlakeU/FS-EEND
import os
import torch
import yaml
import hyperpyyaml
import torch.nn.functional as F
import h5py

from scipy.signal import medfilt
from argparse import ArgumentParser
from pyannote.core import Segment, Timeline, Annotation
from pyannote.metrics.diarization import DiarizationErrorRate
from datasets.dummy_dataset import KaldiDiarizationDataset


def gen_ref(configs, hyp_dir, metric, threshold=0.5, median=11, subsampling=10):
    test_set = KaldiDiarizationDataset(
            data_dir=configs["data"]["val_data_dir"],
            chunk_size=configs["data"]["chunk_size"],
            context_size=configs["data"]["context_recp"],
            input_transform=configs["data"]["feat_type"],
            frame_size=configs["data"]["feat"]["win_length"],
            frame_shift=configs["data"]["feat"]["hop_length"],
            subsampling=configs["data"]["subsampling"],
            rate=configs["data"]["feat"]["sample_rate"],
            label_delay=configs["data"]["label_delay"],
            n_speakers=configs["data"]["num_speakers"],
            use_last_samples=configs["data"]["use_last_samples"],
            shuffle=configs["data"]["shuffle"])
    der = 0.0
    diaerr = 0.0
    spkcon = 0.0
    falarm = 0.0
    miss = 0.0
    spkscore = 0.0
    for i in range(len(test_set)):
        # ref: (5000, C)
        label, rec = test_set.__getfulllabel__(i)

        # perm label
        T, n_spk = label.shape
        frame_idx = torch.arange(1, T + 1).unsqueeze(-1)
        label_idx = (frame_idx * label)
        label_idx = label_idx.masked_fill_(label_idx == 0, torch.inf)
        sort_idx = torch.argsort(torch.min(label_idx, dim=0)[0])
        label = label[:, sort_idx]

        reference = Annotation(uri=f'file{i+1}')
        
        for spkid, frames in enumerate(label.T):
            frames = F.pad(frames, (1, 1), 'constant')
            changes, = torch.where(torch.diff(frames, dim=0) != 0)
            for s, e in zip(changes[::2], changes[1::2]):
                reference[Segment(s, e)] = str(spkid)
        
        # read hypothesis h5 file
        filepath = os.path.join(hyp_dir, rec+".h5")
        data = h5py.File(filepath, 'r')
        pred = torch.where(torch.from_numpy(data['T_hat'][:]).float() > threshold, 1, 0)
        if median > 1:
            pred = medfilt(pred, (median, 1))
        
            pred = torch.from_numpy(pred).float()
        hypothesis =  Annotation(uri=f'file{i+1}')
        for spkid, frames in enumerate(pred.T):
            frames = F.pad(frames, (1, 1), 'constant')
            changes, = torch.where(torch.diff(frames, dim=0) != 0)
            for s, e in zip(changes[::2], changes[1::2]):
                hypothesis[Segment(s*subsampling, e*subsampling)] = str(spkid)
        
        res = metric(reference, hypothesis, detailed=True)
        spkscore += res['total']
        spkcon += res['confusion']
        falarm += res['false alarm']
        miss += res['missed detection']
        diaerr += res['confusion'] + res['false alarm'] + res['missed detection']
        print(rec)
        print("der: ", (res['confusion'] + res['false alarm'] + res['missed detection']) / res['total'])

    der = diaerr / spkscore
    
    print("speaker score: ", spkscore)
    print('mean der: ', der)
    print('mean speaker confusion rate: ', spkcon / spkscore)
    print('mean speaker false alarm rate: ', falarm / spkscore)
    print('mean speaker miss rate: ', miss / spkscore)

            

if __name__ == "__main__":

    parser = ArgumentParser()

    parser.add_argument('--configs', help='Configuration file path', required=True)
    parser.add_argument('--preds_dir', default=None, help='Hypothesis results dir')
    parser.add_argument("--thredshold", default=0.5, help="Threshold of decision")
    parser.add_argument("--median", default=11, help="Median filter parameter")
    setup = parser.parse_args()
    with open(setup.configs, "r") as f:
        configs = hyperpyyaml.load_hyperpyyaml(f)
        f.close()
    preds_dir = "/mnt/home/liangdi/projects/pl_version/pl_eend/tsne_visual/data/onl_2spk_version_tfm_10w_ver_37/preds_h5"

    metric = DiarizationErrorRate(collar=50)
    gen_ref(configs, preds_dir, metric)
metrics.py
Audio-WestlakeU/FS-EEND
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.nn import TransformerEncoder, TransformerEncoderLayer, TransformerDecoder, TransformerDecoderLayer
from ..modules.merge_tfm_encoder import TransformerEncoderFusionLayer
from torch import Tensor
from typing import Optional, Any, Union, Callable

class OnlineTransformerDADiarization(nn.Module):
    def __init__(self, n_speakers, in_size, n_units, n_heads, enc_n_layers, dec_n_layers, dropout, has_mask, max_seqlen, dec_dim_feedforward, conv_delay=9, mask_delay=0, decom_kernel_size=64):
        """ Self-attention-based diarization model.

        Args:
          in_size (int): Dimension of input feature vector
          n_units (int): Number of units in a self-attention block
          n_heads (int): Number of attention heads
          n_layers (int): Number of transformer-encoder layers
          dropout (float): dropout ratio
        """
        super(OnlineTransformerDADiarization, self).__init__()
        self.n_speakers = n_speakers
        self.delay = conv_delay
        self.enc = MaskedTransformerEncoderModel(
            in_size, n_heads, n_units, enc_n_layers, dropout=dropout, has_mask=has_mask, max_seqlen=max_seqlen, mask_delay=mask_delay
        )
        self.dec = MaskedTransformerDecoderModel(
            in_size, n_heads, n_units, dec_n_layers, dim_feedforward=dec_dim_feedforward, dropout=dropout, has_mask=has_mask, max_seqlen=max_seqlen, mask_delay=mask_delay
        )
        self.cnn = nn.Conv1d(n_units, n_units, kernel_size=2 * conv_delay + 1, padding=9)

    def forward(self, src, tgt, ilens):
        n_speakers = [t.shape[1] for t in tgt]
        max_nspks = max(n_speakers)
        # emb: (B, T, E)
        emb = self.enc(src)
        B, T, D = emb.shape
        emb  = [e[:ilen] for e, ilen in zip(emb, ilens)]
        emb = nn.utils.rnn.pad_sequence(emb, padding_value=0, batch_first=True)
        emb: Tensor = self.cnn(emb.transpose(1,2)).transpose(1,2) # (B, T, D)
        emb = emb / torch.norm(emb, dim=-1, keepdim=True)
        attractors = self.dec(emb, max_nspks)
        attractors = attractors / torch.norm(attractors, dim=-1, keepdim=True)

        # Calculate emb consistency loss (cosine similarity)
        attn_map = emb.matmul(emb.transpose(-1, -2))
        # att_norm: (B, T, 1)
        attn_norm = torch.norm(emb, dim=-1, keepdim=True)
        attn_norm = attn_norm.matmul(attn_norm.transpose(-1, -2))
        attn_map = attn_map / (attn_norm + 1e-6)
        tgt_pad = [F.pad(t, (0, max_nspks-t.shape[1]), "constant", 0) for t in tgt]
        tgt_pad = nn.utils.rnn.pad_sequence(tgt_pad, padding_value=0, batch_first=True)
        label_map = tgt_pad.matmul(tgt_pad.transpose(-1, -2))
        tgt_norm = torch.norm(tgt_pad, dim=-1, keepdim=True)
        tgt_norm = tgt_norm.matmul(tgt_norm.transpose(-1, -2))
        label_map = label_map / (tgt_norm + 1e-6)
        emb_consis_loss = F.mse_loss(attn_map, label_map)

        # output: (B, T, C)
        output = torch.matmul(emb.unsqueeze(dim=-2), attractors.transpose(-1, -2)).squeeze(dim=-2)
        output = [out[:ilen, :n_spk] for out, ilen, n_spk in zip(output, ilens, n_speakers)]

        emb = [e[:ilen] for e, ilen in zip(emb, ilens)]
        attractors = [attr[:ilen, 1:n_spk] for attr, ilen, n_spk in zip(attractors, ilens, n_speakers)]
        return output, emb_consis_loss, emb, attractors
    
    def test(self, src, ilens, max_nspks=6):
        # emb: (B, T, E)
        emb = self.enc(src)
        B, T, D = emb.shape
        emb  = [e[:ilen] for e, ilen in zip(emb, ilens)]
        emb = nn.utils.rnn.pad_sequence(emb, padding_value=0, batch_first=True)
        emb: Tensor = self.cnn(emb.transpose(1,2)).transpose(1,2) # (B, T, D)
        emb = emb / torch.norm(emb, dim=-1, keepdim=True)
        attractors = self.dec(emb, max_nspks)
        attractors = attractors / torch.norm(attractors, dim=-1, keepdim=True)
        
        # output: (B, T, C)
        output = torch.matmul(emb.unsqueeze(dim=-2), attractors.transpose(-1, -2)).squeeze(dim=-2)
        output = [out[:ilen] for out, ilen in zip(output, ilens)]

        emb = [e[:ilen] for e, ilen in zip(emb, ilens)]
        attractors = [attr[:ilen] for attr, ilen in zip(attractors, ilens)]
        return output, emb, attractors


class MaskedTransformerDecoderModel(nn.Module):
    def __init__(self, in_size, n_heads, n_units, n_layers, dim_feedforward, dropout=0.5, has_mask=False, max_seqlen = 500, has_pos=False, mask_delay=0):
        super(MaskedTransformerDecoderModel, self).__init__()
        self.in_size = in_size
        self.n_heads = n_heads
        self.n_units = n_units
        self.n_layers = n_layers
        self.has_pos = has_pos
        self.has_mask = has_mask
        self.max_seqlen = max_seqlen
        self.mask_delay = mask_delay

        self.encoder = nn.Linear(in_size, n_units)
        self.encoder_norm = nn.LayerNorm(n_units)

        self.pos_enc = PositionalEncoding(n_units, dropout)
        self.convert = nn.Linear(n_units * 2, n_units)
        decoder_layers = TransformerEncoderFusionLayer(n_units, n_heads, dim_feedforward, dropout, batch_first=True)
        self.attractor_decoder = TransformerEncoder(decoder_layers, n_layers)

    def _generate_square_subsequent_mask(self, sz, device):
        mask = (torch.triu(torch.ones(sz, sz, device=device), diagonal=-self.mask_delay) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def forward(self, emb: Tensor, max_nspks: int, activation: Optional[Callable]=None):
        pos_enc = self.pos_enc(emb, max_nspks) # (B, T, C, D)
        attractors_init: Tensor = self.convert(torch.cat([emb.unsqueeze(dim=2).repeat(1, 1, max_nspks, 1), pos_enc], dim=-1))
        
        t_mask = self._generate_square_subsequent_mask(emb.shape[1], emb.device)
        attractors = self.attractor_decoder(attractors_init, t_mask)
        return attractors

class MaskedTransformerEncoderModel(nn.Module):
    def __init__(self, in_size, n_heads, n_units, n_layers, dim_feedforward=2048, dropout=0.5, has_mask=False, max_seqlen = 500, has_pos=False, mask_delay=0):
        """ Self-attention-based diarization model.

        Args:
          n_speakers (int): Number of speakers in recording
          in_size (int): Dimension of input feature vector
          n_heads (int): Number of attention heads
          n_units (int): Number of units in a self-attention block
          n_layers (int): Number of transformer-encoder layers
          dropout (float): dropout ratio
        """
        super(MaskedTransformerEncoderModel, self).__init__()
        self.in_size = in_size
        self.n_heads = n_heads
        self.n_units = n_units
        self.n_layers = n_layers
        self.has_pos = has_pos
        self.has_mask = has_mask
        self.max_seqlen = max_seqlen
        self.mask_delay = mask_delay

        self.encoder = nn.Linear(in_size, n_units)
        self.encoder_norm = nn.LayerNorm(n_units)
        if self.has_pos:
            self.pos_encoder = PositionalEncoding(n_units, dropout)
        encoder_layers = TransformerEncoderLayer(n_units, n_heads, dim_feedforward, dropout)
        self.transformer_encoder = TransformerEncoder(encoder_layers, n_layers)

        self.init_weights()

    def _generate_square_subsequent_mask(self, sz, device):
        mask = (torch.triu(torch.ones(sz, sz, device=device), diagonal=-self.mask_delay) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def init_weights(self):
        initrange = 0.1
        self.encoder.bias.data.zero_()
        self.encoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, src, activation=None):

        ilens = [x.shape[0] for x in src]
        src = nn.utils.rnn.pad_sequence(src, padding_value=-1, batch_first=True)

        src_mask = None
        if self.has_mask:
            src_mask = self._generate_square_subsequent_mask(src.shape[1], src.device)

        # src: (B, T, E)
        src = self.encoder(src)
        src = self.encoder_norm(src)
        # src: (T, B, E)
        src = src.transpose(0, 1)
        if self.has_pos:
            # src: (T, B, E)
            src = self.pos_encoder(src)
        # output: (T, B, E)
        output = self.transformer_encoder(src, src_mask)
        # output: (B, T, E)
        output = output.transpose(0, 1)

        if activation:
            output = activation(output)

        return output

class PositionalEncoding(nn.Module):
    """Inject some information about the relative or absolute position of the tokens
        in the sequence. The positional encodings have the same dimension as
        the embeddings, so that the two can be summed. Here, we use sine and cosine
        functions of different frequencies.
    .. math::
        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
        \text{where pos is the word position and i is the embed idx)
    Args:
        d_model: the embed dim (required).
        dropout: the dropout value (default=0.1).
        max_len: the max. length of the incoming sequence (default=5000).
    Examples:
        >>> pos_encoder = PositionalEncoding(d_model)
    """
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x: Tensor, max_nspks):
        # Add positional information to each time step of x
        pe = self.pe[:, :max_nspks, :]
        pe = pe.unsqueeze(dim=0).repeat(x.shape[0], x.shape[1], 1, 1) # (B, T, C, D)
        x = x.unsqueeze(dim=2).repeat(1, 1, max_nspks, 1)
        # x = x + pe
        return pe
nnet/model/onl_tfm_enc_1dcnn_enc_linear_non_autoreg_pos_enc_l2norm.py
Audio-WestlakeU/FS-EEND
import copy
from typing import Optional, Any, Union, Callable

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn.modules import Module
from torch.nn.modules.activation import MultiheadAttention
from torch.nn.modules.container import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn.modules.dropout import Dropout
from torch.nn.modules.linear import Linear
from torch.nn.modules.normalization import LayerNorm

class TransformerEncoderFusionLayer(Module):
    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
    This standard encoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of the intermediate layer, can be a string
            ("relu" or "gelu") or a unary callable. Default: relu
        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
        batch_first: If ``True``, then the input and output tensors are provided
            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
        norm_first: if ``True``, layer norm is done prior to attention and feedforward
            operations, respectively. Otherwise it's done after. Default: ``False`` (after).

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> src = torch.rand(10, 32, 512)
        >>> out = encoder_layer(src)

    Alternatively, when ``batch_first`` is ``True``:
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
        >>> src = torch.rand(32, 10, 512)
        >>> out = encoder_layer(src)

    Fast path:
        forward() will use a special optimized implementation if all of the following
        conditions are met:

        - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor
          argument ``requires_grad``
        - training is disabled (using ``.eval()``)
        - batch_first is ``True`` and the input is batched (i.e., ``src.dim() == 3``)
        - activation is one of: ``"relu"``, ``"gelu"``, ``torch.functional.relu``, or ``torch.functional.gelu``
        - at most one of ``src_mask`` and ``src_key_padding_mask`` is passed
        - if src is a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_, neither ``src_mask``
          nor ``src_key_padding_mask`` is passed
        - the two ``LayerNorm`` instances have a consistent ``eps`` value (this will naturally be the case
          unless the caller has manually modified one without modifying the other)

        If the optimized implementation is in use, a
        `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be
        passed for ``src`` to represent padding more efficiently than using a padding
        mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ will be
        returned, and an additional speedup proportional to the fraction of the input that
        is padding can be expected.
    """
    __constants__ = ['batch_first', 'norm_first']

    def __init__(self, d_model: int, nhead: int, dim_feedforward: int = 2048, dropout: float = 0.1,
                 activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
                 layer_norm_eps: float = 1e-5, batch_first: bool = False, norm_first: bool = False,
                 device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(TransformerEncoderFusionLayer, self).__init__()
        self.self_attn1 = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
                                            **factory_kwargs)
        self.self_attn2 = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
                                            **factory_kwargs)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)

        self.norm_first = norm_first
        self.norm11 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.norm12 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.norm21 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.norm22 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.dropout11 = Dropout(dropout)
        self.dropout21 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        # Legacy string support for activation function.
        if isinstance(activation, str):
            activation = _get_activation_fn(activation)

        # We can't test self.activation in forward() in TorchScript,
        # so stash some information about it instead.
        if activation is F.relu or isinstance(activation, torch.nn.ReLU):
            self.activation_relu_or_gelu = 1
        elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
            self.activation_relu_or_gelu = 2
        else:
            self.activation_relu_or_gelu = 0
        self.activation = activation

    def __setstate__(self, state):
        super(TransformerEncoderFusionLayer, self).__setstate__(state)
        if not hasattr(self, 'activation'):
            self.activation = F.relu


    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the input through the encoder layer.

        Args:
            src: the sequence to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """

        if src_key_padding_mask is not None:
            _skpm_dtype = src_key_padding_mask.dtype
            if _skpm_dtype != torch.bool and not torch.is_floating_point(src_key_padding_mask):
                raise AssertionError(
                    "only bool and floating types of key_padding_mask are supported")
        # see Fig. 1 of https://arxiv.org/pdf/2002.04745v1.pdf
        why_not_sparsity_fast_path = ''
        if not src.dim() == 3:
            why_not_sparsity_fast_path = f"input not batched; expected src.dim() of 3 but got {src.dim()}"
        elif self.training:
            why_not_sparsity_fast_path = "training is enabled"
        elif not self.self_attn1.batch_first :
            why_not_sparsity_fast_path = "self_attn.batch_first was not True"
        elif not self.self_attn1._qkv_same_embed_dim :
            why_not_sparsity_fast_path = "self_attn._qkv_same_embed_dim was not True"
        elif not self.activation_relu_or_gelu:
            why_not_sparsity_fast_path = "activation_relu_or_gelu was not True"
        elif not (self.norm11.eps == self.norm12.eps):
            why_not_sparsity_fast_path = "norm1.eps is not equal to norm2.eps"
        elif src_mask is not None:
            why_not_sparsity_fast_path = "src_mask is not supported for fastpath"
        elif src.is_nested and src_key_padding_mask is not None:
            why_not_sparsity_fast_path = "src_key_padding_mask is not supported with NestedTensor input for fastpath"
        elif self.self_attn1.num_heads % 2 == 1:
            why_not_sparsity_fast_path = "num_head is odd"
        elif torch.is_autocast_enabled():
            why_not_sparsity_fast_path = "autocast is enabled"

        if not why_not_sparsity_fast_path:
            tensor_args = (
                src,
                self.self_attn1.in_proj_weight,
                self.self_attn1.in_proj_bias,
                self.self_attn1.out_proj.weight,
                self.self_attn1.out_proj.bias,
                self.self_attn2.in_proj_weight,
                self.self_attn2.in_proj_bias,
                self.self_attn2.out_proj.weight,
                self.self_attn2.out_proj.bias,
                self.norm11.weight,
                self.norm11.bias,
                self.norm12.weight,
                self.norm12.bias,
                self.norm21.weight,
                self.norm21.bias,
                self.norm22.weight,
                self.norm22.bias,
                self.linear1.weight,
                self.linear1.bias,
                self.linear2.weight,
                self.linear2.bias,
            )

            # We have to use list comprehensions below because TorchScript does not support
            # generator expressions.
            if torch.overrides.has_torch_function(tensor_args):
                why_not_sparsity_fast_path = "some Tensor argument has_torch_function"
            elif not all((x.is_cuda or 'cpu' in str(x.device)) for x in tensor_args):
                why_not_sparsity_fast_path = "some Tensor argument is neither CUDA nor CPU"
            elif torch.is_grad_enabled() and any(x.requires_grad for x in tensor_args):
                why_not_sparsity_fast_path = ("grad is enabled and at least one of query or the "
                                              "input/output projection weights or biases requires_grad")

            if not why_not_sparsity_fast_path:
                return torch._transformer_encoder_layer_fwd(
                    src,
                    self.self_attn1.embed_dim,
                    self.self_attn1.num_heads,
                    self.self_attn1.in_proj_weight,
                    self.self_attn1.in_proj_bias,
                    self.self_attn1.out_proj.weight,
                    self.self_attn1.out_proj.bias,
                    self.self_attn2.embed_dim,
                    self.self_attn2.num_heads,
                    self.self_attn2.in_proj_weight,
                    self.self_attn2.in_proj_bias,
                    self.self_attn2.out_proj.weight,
                    self.self_attn2.out_proj.bias,
                    self.activation_relu_or_gelu == 2,
                    self.norm_first,
                    self.norm11.eps,
                    self.norm11.weight,
                    self.norm11.bias,
                    self.norm12.weight,
                    self.norm12.bias,
                    self.norm21.eps,
                    self.norm21.weight,
                    self.norm21.bias,
                    self.norm22.weight,
                    self.norm22.bias,
                    self.linear1.weight,
                    self.linear1.bias,
                    self.linear2.weight,
                    self.linear2.bias,
                    # TODO: if src_mask and src_key_padding_mask merge to single 4-dim mask
                    src_mask if src_mask is not None else src_key_padding_mask,
                    1 if src_key_padding_mask is not None else
                    0 if src_mask is not None else
                    None,
                )


        # Self-attn on time-frame dim
        B, T, C, D = src.shape
        x = src.transpose(1, 2).reshape(B*C, T, D)
        if self.norm_first:
            x = x + self._sa_block1(self.norm11(x), src_mask, src_key_padding_mask)
            # x = x + self._ff_block(self.norm12(x))
        else:
            x = self.norm11(x + self._sa_block1(x, src_mask, src_key_padding_mask))
            # x = self.norm12(x + self._ff_block(x))
        x = x.reshape(B, C, T, D).transpose(1, 2).reshape(B*T, C, D)
        
        # Self-attention on spk dim
        if self.norm_first:
            x = x + self._sa_block2(self.norm21(x), None, None)
            x = x + self._ff_block(self.norm22(x))
        else:
            x = self.norm21(x + self._sa_block2(x, None, None))
            x = self.norm22(x + self._ff_block(x))
        x = x.reshape(B, T, C, D)

        return x

    # self-attention block
    def _sa_block1(self, x: Tensor,
                  attn_mask: Optional[Tensor], key_padding_mask: Optional[Tensor]) -> Tensor:
        x = self.self_attn1(x, x, x,
                           attn_mask=attn_mask,
                           key_padding_mask=key_padding_mask,
                           need_weights=False)[0]
        return self.dropout11(x)
    
    # self-attention2 block
    def _sa_block2(self, x: Tensor,
                  attn_mask: Optional[Tensor], key_padding_mask: Optional[Tensor]) -> Tensor:
        x = self.self_attn2(x, x, x,
                           attn_mask=attn_mask,
                           key_padding_mask=key_padding_mask,
                           need_weights=False)[0]
        return self.dropout21(x)

    # feed forward block
    def _ff_block(self, x: Tensor) -> Tensor:
        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
        return self.dropout2(x)


class TransformerEncoderLayerC(Module):
    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
    This standard encoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of the intermediate layer, can be a string
            ("relu" or "gelu") or a unary callable. Default: relu
        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
        batch_first: If ``True``, then the input and output tensors are provided
            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
        norm_first: if ``True``, layer norm is done prior to attention and feedforward
            operations, respectively. Otherwise it's done after. Default: ``False`` (after).

    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
        >>> src = torch.rand(10, 32, 512)
        >>> out = encoder_layer(src)

    Alternatively, when ``batch_first`` is ``True``:
        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
        >>> src = torch.rand(32, 10, 512)
        >>> out = encoder_layer(src)

    Fast path:
        forward() will use a special optimized implementation if all of the following
        conditions are met:

        - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor
          argument ``requires_grad``
        - training is disabled (using ``.eval()``)
        - batch_first is ``True`` and the input is batched (i.e., ``src.dim() == 3``)
        - activation is one of: ``"relu"``, ``"gelu"``, ``torch.functional.relu``, or ``torch.functional.gelu``
        - at most one of ``src_mask`` and ``src_key_padding_mask`` is passed
        - if src is a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_, neither ``src_mask``
          nor ``src_key_padding_mask`` is passed
        - the two ``LayerNorm`` instances have a consistent ``eps`` value (this will naturally be the case
          unless the caller has manually modified one without modifying the other)

        If the optimized implementation is in use, a
        `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be
        passed for ``src`` to represent padding more efficiently than using a padding
        mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ will be
        returned, and an additional speedup proportional to the fraction of the input that
        is padding can be expected.
    """
    __constants__ = ['batch_first', 'norm_first']

    def __init__(self, d_model: int, nhead: int, dim_feedforward: int = 2048, dropout: float = 0.1,
                 activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
                 layer_norm_eps: float = 1e-5, batch_first: bool = False, norm_first: bool = False,
                 device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(TransformerEncoderLayerC, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
                                            **factory_kwargs)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)

        self.norm_first = norm_first
        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        # Legacy string support for activation function.
        if isinstance(activation, str):
            activation = _get_activation_fn(activation)

        # We can't test self.activation in forward() in TorchScript,
        # so stash some information about it instead.
        if activation is F.relu or isinstance(activation, torch.nn.ReLU):
            self.activation_relu_or_gelu = 1
        elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
            self.activation_relu_or_gelu = 2
        else:
            self.activation_relu_or_gelu = 0
        self.activation = activation

    def __setstate__(self, state):
        super(TransformerEncoderLayerC, self).__setstate__(state)
        if not hasattr(self, 'activation'):
            self.activation = F.relu


    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
        r"""Pass the input through the encoder layer.

        Args:
            src: the sequence to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """

        if src_key_padding_mask is not None:
            _skpm_dtype = src_key_padding_mask.dtype
            if _skpm_dtype != torch.bool and not torch.is_floating_point(src_key_padding_mask):
                raise AssertionError(
                    "only bool and floating types of key_padding_mask are supported")
        # see Fig. 1 of https://arxiv.org/pdf/2002.04745v1.pdf
        why_not_sparsity_fast_path = ''
        if not src.dim() == 3:
            why_not_sparsity_fast_path = f"input not batched; expected src.dim() of 3 but got {src.dim()}"
        elif self.training:
            why_not_sparsity_fast_path = "training is enabled"
        elif not self.self_attn.batch_first :
            why_not_sparsity_fast_path = "self_attn.batch_first was not True"
        elif not self.self_attn._qkv_same_embed_dim :
            why_not_sparsity_fast_path = "self_attn._qkv_same_embed_dim was not True"
        elif not self.activation_relu_or_gelu:
            why_not_sparsity_fast_path = "activation_relu_or_gelu was not True"
        elif not (self.norm1.eps == self.norm2.eps):
            why_not_sparsity_fast_path = "norm1.eps is not equal to norm2.eps"
        elif src_mask is not None:
            why_not_sparsity_fast_path = "src_mask is not supported for fastpath"
        elif src.is_nested and src_key_padding_mask is not None:
            why_not_sparsity_fast_path = "src_key_padding_mask is not supported with NestedTensor input for fastpath"
        elif self.self_attn.num_heads % 2 == 1:
            why_not_sparsity_fast_path = "num_head is odd"
        elif torch.is_autocast_enabled():
            why_not_sparsity_fast_path = "autocast is enabled"

        if not why_not_sparsity_fast_path:
            tensor_args = (
                src,
                self.self_attn.in_proj_weight,
                self.self_attn.in_proj_bias,
                self.self_attn.out_proj.weight,
                self.self_attn.out_proj.bias,
                self.norm1.weight,
                self.norm1.bias,
                self.norm2.weight,
                self.norm2.bias,
                self.linear1.weight,
                self.linear1.bias,
                self.linear2.weight,
                self.linear2.bias,
            )

            # We have to use list comprehensions below because TorchScript does not support
            # generator expressions.
            if torch.overrides.has_torch_function(tensor_args):
                why_not_sparsity_fast_path = "some Tensor argument has_torch_function"
            elif not all((x.is_cuda or 'cpu' in str(x.device)) for x in tensor_args):
                why_not_sparsity_fast_path = "some Tensor argument is neither CUDA nor CPU"
            elif torch.is_grad_enabled() and any(x.requires_grad for x in tensor_args):
                why_not_sparsity_fast_path = ("grad is enabled and at least one of query or the "
                                              "input/output projection weights or biases requires_grad")

            if not why_not_sparsity_fast_path:
                return torch._transformer_encoder_layer_fwd(
                    src,
                    self.self_attn.embed_dim,
                    self.self_attn.num_heads,
                    self.self_attn.in_proj_weight,
                    self.self_attn.in_proj_bias,
                    self.self_attn.out_proj.weight,
                    self.self_attn.out_proj.bias,
                    self.activation_relu_or_gelu == 2,
                    self.norm_first,
                    self.norm1.eps,
                    self.norm1.weight,
                    self.norm1.bias,
                    self.norm2.weight,
                    self.norm2.bias,
                    self.linear1.weight,
                    self.linear1.bias,
                    self.linear2.weight,
                    self.linear2.bias,
                    # TODO: if src_mask and src_key_padding_mask merge to single 4-dim mask
                    src_mask if src_mask is not None else src_key_padding_mask,
                    1 if src_key_padding_mask is not None else
                    0 if src_mask is not None else
                    None,
                )


        x = src
        if self.norm_first:
            x = x + self._sa_block(self.norm1(x), src_mask, src_key_padding_mask)
            x = x + self._ff_block(self.norm2(x))
        else:
            x = self.norm1(x + self._sa_block(x, src_mask, src_key_padding_mask))
            x = self.norm2(x + self._ff_block(x))

        return x

    # self-attention block
    def _sa_block(self, x: Tensor,
                  attn_mask: Optional[Tensor], key_padding_mask: Optional[Tensor]) -> Tensor:
        x = self.self_attn(x, x, x,
                           attn_mask=attn_mask,
                           key_padding_mask=key_padding_mask,
                           need_weights=False)[0]
        return self.dropout1(x)

    # feed forward block
    def _ff_block(self, x: Tensor) -> Tensor:
        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
        return self.dropout2(x)



def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu

    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
nnet/modules/merge_tfm_encoder.py
Audio-WestlakeU/FS-EEND
import os
import torch
import torch.nn.functional as F
import numpy as np
import pytorch_lightning as pl

from collections import defaultdict
from pprint import pprint
from torch.utils.data import DataLoader
from torchaudio.transforms import MelSpectrogram, AmplitudeToDB
from .utils.loss import batch_pit_n_speaker_loss, report_diarization_error, standard_loss
from .utils.utils import ContextBuilder, TorchScaler


import time

class SpeakerDiarization(pl.LightningModule):
    """
    Speaker diarization main stream: 
    TODO: 
        test step is currently the same as the validation step
    """
    def __init__(self, hparams, model, datasets, opt, scheduler, collate_func) -> None:
        super(SpeakerDiarization, self).__init__()
        self.hparams.update(hparams)
        # Training configs
        self.datasets = datasets
        self.model = model
        self.opt = opt
        self.collate_func = collate_func
        self.scheduler = scheduler
        # self.loss_func = PITLoss(self.hparams["data"]["num_speakers"])
        self.loss_func1 = batch_pit_n_speaker_loss
        self.loss_func2 = standard_loss
        self.max_spks = self.hparams["data"]["max_speakers"]
        self.label_delay = self.hparams["data"]["label_delay"]

    def scaler_init(self):
        return TorchScaler(
                self.hparams["data"]["scaler"]["statistic"],
                self.hparams["data"]["scaler"]["normtype"],
                self.hparams["data"]["scaler"]["dims"],
            )

    
    def detect(self, x, labels, ilens):
        # Here we could apply normalization or something else
        return self.model(x, tgt=labels, ilens=ilens)
    
    def training_step(self, batch, batch_index):
        feats, labels, _ = batch
        clip_lengths = [x.shape[0] for x in feats]
        n_spks = [l.shape[1] for l in labels]
        max_spk = max(n_spks)
        labels = [F.pad(l, (0,max_spk-l.shape[1]), "constant", 0) for l in labels]

        # perm labels
        labels = torch.nn.utils.rnn.pad_sequence(labels, padding_value=0, batch_first=True)
        B, T, max_spk = labels.shape
        frame_index = torch.arange(1, T + 1, device=feats[0].device).unsqueeze(0).unsqueeze(-1)
        labels_index = (frame_index * labels)
        labels_index = labels_index.masked_fill_(labels_index == 0, torch.inf)

        # [B, max_spk]
        sort_index = torch.argsort(torch.min(labels_index, dim=1)[0], dim=1)
        labels = labels[torch.arange(B).unsqueeze(1), :, sort_index].transpose(-1, -2)

        labels_silence = torch.ones([labels.shape[0], labels.shape[1]], device=labels.device) - labels.max(-1)[0]
        labels = torch.cat([labels_silence.unsqueeze(-1), labels], dim=-1)

        labels_nonespk = torch.zeros([B, T, 1], device=labels.device)
        labels = torch.cat([labels, labels_nonespk], dim=-1)

        labels = [l[:ilen, :nspk+2] for l, ilen, nspk in zip(labels, clip_lengths, n_spks)]

        # Transform feats to MelSpectrum
        preds, emb_loss, embs, attractors = self.detect(feats, labels, clip_lengths)
        # print(preds.shape)

        # pit_loss1, perm_labels = self.loss_func1(preds, labels)
        pit_loss = self.loss_func2(preds, labels, label_delay=self.label_delay)

        # Calculate the total loss
        tot_loss = pit_loss + emb_loss
        self.log("train/lr", self.opt.param_groups[-1]["lr"], prog_bar=True)
        self.log("train/pit_loss", pit_loss)
        self.log("train/emb_loss", emb_loss)
        self.log("train/tot_loss", tot_loss)

        return tot_loss

    def validation_step(self, batch, batch_index):
        feats, labels, _ = batch
        # Model pred
        clip_lengths = [x.shape[0] for x in feats]
        n_spks = [l.shape[1] for l in labels]
        max_spk = max(n_spks)
        labels = [F.pad(l, (0,max_spk-l.shape[1]), "constant", 0) for l in labels]

        # perm labels
        labels = torch.nn.utils.rnn.pad_sequence(labels, padding_value=0, batch_first=True)
        B, T, max_spk = labels.shape
        frame_index = torch.arange(1, T + 1, device=feats[0].device).unsqueeze(0).unsqueeze(-1)
        labels_index = (frame_index * labels)
        labels_index = labels_index.masked_fill_(labels_index == 0, torch.inf)

        # [B, max_spk]
        sort_index = torch.argsort(torch.min(labels_index, dim=1)[0], dim=1)
        labels = labels[torch.arange(B).unsqueeze(1), :, sort_index].transpose(-1, -2)

        labels_silence = torch.ones([labels.shape[0], labels.shape[1]], device=labels.device) - labels.max(-1)[0]
        labels = torch.cat([labels_silence.unsqueeze(-1), labels], dim=-1)
        
        labels_nonespk = torch.zeros([B, T, 1], device=labels.device)
        labels = torch.cat([labels, labels_nonespk], dim=-1)

        labels = [l[:ilen, :nspk+2] for l, ilen, nspk in zip(labels, clip_lengths, n_spks)]

        preds, emb_loss, embs, attractors = self.detect(feats, labels, clip_lengths)
        #pit_loss1, perm_labels = self.loss_func1(preds, labels)
        pit_loss = self.loss_func2(preds, labels, label_delay=self.label_delay)
        # print(f"batch_lost_n_speakers: {pit_loss1},  batch_pit_loss: {pit_loss3},  pit_loss: {pit_loss}")
        tot_loss = pit_loss + emb_loss

        # Metrics
        preds_realspk = [p[:, 1:-1] for p in preds]
        labels_realspk = [l[:, 1:-1] for l in labels]
        stats = report_diarization_error(preds_realspk, labels_realspk, label_delay=self.label_delay)

        self.log("val/pit_loss", pit_loss)
        self.log("val/emb_loss", emb_loss)
        self.log("val/tot_loss", tot_loss)

        return stats
    
    def validation_epoch_end(self, val_step_outputs) -> None:
        stats_holder = defaultdict(list)
        for stats in val_step_outputs:
            for k, v in stats.items():
                stats_holder[k] += v
        # Calculate DER
        stats_avg = {k: sum(v)/len(v) for k, v in stats_holder.items()}
        stats_avg["DER"] = stats_avg["diarization_error"] / stats_avg["speaker_scored"]
        obj_metric = stats_avg['DER']
        self.log("val/obj_metric", obj_metric, sync_dist=True)
        self.log("val/DER", stats_avg['DER'], sync_dist=True)
        self.log("val/speech_scored", stats_avg['speech_scored'], sync_dist=True)
        self.log("val/speech_miss", stats_avg['speech_miss'], sync_dist=True)
        self.log("val/speech_falarm", stats_avg['speech_falarm'], sync_dist=True)
        self.log("val/speaker_scored", stats_avg['speaker_scored'], sync_dist=True)
        self.log("val/speaker_miss", stats_avg['speaker_miss'], sync_dist=True)
        self.log("val/speaker_falarm", stats_avg['speaker_falarm'], sync_dist=True)
        self.log("val/speaker_error", stats_avg['speaker_error'], sync_dist=True)
        self.log("val/diarization_error", stats_avg['diarization_error'], sync_dist=True)
        self.log("val/frames", stats_avg['frames'], sync_dist=True)


    def test_step(self, batch, batch_index):
        feats, labels, rec = batch
        # Model pred
        clip_lengths = [x.shape[0] for x in feats]
        n_spks = [l.shape[1] for l in labels]
        max_spk = max(n_spks)
        labels = [F.pad(l, (0,max_spk-l.shape[1]), "constant", 0) for l in labels]

        # perm labels
        labels = torch.nn.utils.rnn.pad_sequence(labels, padding_value=0, batch_first=True)
        B, T, max_spk = labels.shape
        frame_index = torch.arange(1, T + 1, device=feats[0].device).unsqueeze(0).unsqueeze(-1)
        labels_index = (frame_index * labels)
        labels_index = labels_index.masked_fill_(labels_index == 0, torch.inf)

        # [B, max_spk]
        sort_index = torch.argsort(torch.min(labels_index, dim=1)[0], dim=1)
        labels = labels[torch.arange(B).unsqueeze(1), :, sort_index].transpose(-1, -2)

        labels_silence = torch.ones([labels.shape[0], labels.shape[1]], device=labels.device) - labels.max(-1)[0]
        labels = torch.cat([labels_silence.unsqueeze(-1), labels], dim=-1)

        labels_nonespk = torch.zeros([B, T, 1], device=labels.device)
        labels = torch.cat([labels, labels_nonespk], dim=-1)

        labels = [l[:ilen, :nspk+2] for l, ilen, nspk in zip(labels, clip_lengths, n_spks)]
        
        preds, embs, attractors = self.model.test(feats, clip_lengths, self.max_spks + 2)
        # pit_loss, perm_labels = self.loss_func1(preds, labels)
        # pit_loss = self.loss_func2(preds, labels)
        # tot_loss = pit_loss + attr_loss + emb_loss

        # Metrics
        preds_realspk = [p[:, 1:nspk + 1] for p, nspk in zip(preds, n_spks)]
        labels_realspk = [l[:, 1:-1] for l in labels]
        stats = report_diarization_error(preds_realspk, labels_realspk, label_delay=self.label_delay)

        label_delay = self.hparams["data"]["label_delay"]
        preds = [torch.cat([p[label_delay:, 1:], p[-1, 1:].unsqueeze(0).repeat(label_delay, 1)], dim=0)  for p in preds]
        num_spks = "2"
        version = "_tfm_10w_ver_0"
        # save_dir_parnt = f"/mnt/home/liangdi/projects/pl_version/pl_eend_nonautoreg/tsne_visual/data/onl_{num_spks}spk_version{version}"
        # for content in ["preds", "labels", "embs", "attractors"]:
        #     save_dir = os.path.join(save_dir_parnt, content)
        #     if not os.path.isdir(save_dir):
        #         os.makedirs(save_dir)
        # np.save(f"/mnt/home/liangdi/projects/pl_version/pl_eend_nonautoreg/tsne_visual/data/onl_{num_spks}spk_version{version}/preds/{rec[0]}.npy", preds[0].detach().cpu().numpy())
        # np.save(f"/mnt/home/liangdi/projects/pl_version/pl_eend_nonautoreg/tsne_visual/data/onl_{num_spks}spk_version{version}/labels/label_{rec[0]}.npy", labels_realspk[0].detach().cpu().numpy())
        # np.save(f"/mnt/home/liangdi/projects/pl_version/pl_eend_nonautoreg/tsne_visual/data/onl_{num_spks}spk_version{version}/embs/emb_{rec[0]}.npy", embs[0].detach().cpu().numpy())
        # np.save(f"/mnt/home/liangdi/projects/pl_version/pl_eend_nonautoreg/tsne_visual/data/onl_{num_spks}spk_version{version}/attractors/attractor_{rec[0]}.npy", attractors[0].detach().cpu().numpy())
        return stats

    def test_epoch_end(self, test_step_outputs) -> None:
        stats_holder = defaultdict(list)
        for stats in test_step_outputs:
            for k, v in stats.items():
                stats_holder[k] += v
        # [print(x) for x in stats_holder["speaker_scored"] if not x ]
        # Calculate DER
        
        stats_avg = {k: sum(v)/len(v) for k, v in stats_holder.items()}
        stats_avg["DER"] = stats_avg["diarization_error"] / stats_avg["speaker_scored"]
        results = {
            "test/DER": stats_avg['DER'],
            "test/speech_scored": stats_avg['speech_scored'],
            "test/speech_miss": stats_avg['speech_miss'],
            "test/speech_falarm": stats_avg['speech_falarm'],
            "test/speaker_scored": stats_avg['speaker_scored'],
            "test/speaker_miss": stats_avg['speaker_miss'],
            "test/speaker_falarm": stats_avg['speaker_falarm'],
            "test/speaker_error": stats_avg['speaker_error'],
            "test/diarization_error": stats_avg['diarization_error'],
            "test/frames": stats_avg['frames'],
            "test/speaker_miss_rate": stats_avg['speaker_miss'] / stats_avg['speaker_scored'],
            "test/speaker_falarm_rate": stats_avg['speaker_falarm'] / stats_avg['speaker_scored'],
            "test/speaker_error_rate": stats_avg['speaker_error'] / stats_avg['speaker_scored'],
        }
        pprint(results)
        self.logger.log_metrics(results)
        self.logger.log_hyperparams(self.hparams, results)
        
        for key in results.keys():
            self.log(key, results[key], prog_bar=True, logger=False)


    def configure_optimizers(self):
        if self.scheduler is not None:
            return {"optimizer": self.opt, "lr_scheduler": {"scheduler": self.scheduler, "interval": "step"}}
        else:
            return {"optimizer": self.opt}

    def train_dataloader(self):
        return DataLoader(
            self.datasets["train"],
            batch_size=self.hparams["training"]["batch_size"],
            shuffle=self.hparams["training"]["shuffle"],
            num_workers=self.hparams["training"]["n_workers"],
            collate_fn=self.collate_func
        )
    
    def val_dataloader(self):
        return DataLoader(
            self.datasets["val"],
            batch_size=self.hparams["training"]["batch_size"],
            shuffle=False,
            num_workers=self.hparams["training"]["n_workers"],
            collate_fn=self.collate_func
        )
    
    def test_dataloader(self):
        return DataLoader(
            self.datasets["val"],
            batch_size=self.hparams["training"]["batch_size"],
            shuffle=False,
            num_workers=self.hparams["training"]["n_workers"],
            collate_fn=self.collate_func
        )
    

train/oln_tfm_enc_dec.py
Audio-WestlakeU/FS-EEND
import os
import torch
import torch.nn.functional as F
import numpy as np
import pytorch_lightning as pl

from collections import defaultdict
from pprint import pprint
from torch.utils.data import DataLoader
from torchaudio.transforms import MelSpectrogram, AmplitudeToDB
from .utils.loss import batch_pit_n_speaker_loss, report_diarization_error, standard_loss, pad_labels, pad_preds
from .utils.utils import ContextBuilder, TorchScaler

import time

class SpeakerDiarization(pl.LightningModule):
    """
    Speaker diarization main stream: 
    TODO: 
        test step is currently the same as the validation step
    """
    def __init__(self, hparams, model, datasets, opt, scheduler, collate_func) -> None:
        super(SpeakerDiarization, self).__init__()
        self.hparams.update(hparams)
        # Training configs
        self.datasets = datasets
        self.model = model
        self.opt = opt
        self.collate_func = collate_func
        self.scheduler = scheduler
        # self.loss_func = PITLoss(self.hparams["data"]["num_speakers"])
        self.loss_func1 = batch_pit_n_speaker_loss
        self.loss_func2 = standard_loss
        self.max_spks = self.hparams["data"]["max_speakers"]
        self.label_delay = self.hparams["data"]["label_delay"]

    def scaler_init(self):
        return TorchScaler(
                self.hparams["data"]["scaler"]["statistic"],
                self.hparams["data"]["scaler"]["normtype"],
                self.hparams["data"]["scaler"]["dims"],
            )

    
    def detect(self, x, labels, ilens):
        # Here we could apply normalization or something else
        return self.model(x, tgt=labels, ilens=ilens)
    
    def training_step(self, batch, batch_index):
        feats, labels, _ = batch
        clip_lengths = [x.shape[0] for x in feats]
        n_spks = [l.shape[1] for l in labels]
        max_spk = max(n_spks)
        labels = [F.pad(l, (0,max_spk-l.shape[1]), "constant", 0) for l in labels]

        # perm labels
        labels = torch.nn.utils.rnn.pad_sequence(labels, padding_value=0, batch_first=True)
        B, T, max_spk = labels.shape
        frame_index = torch.arange(1, T + 1).to(feats[0]).unsqueeze(0).unsqueeze(-1)
        labels_index = (frame_index * labels)
        labels_index = labels_index.masked_fill_(labels_index == 0, torch.inf)

        # [B, max_spk]
        sort_index = torch.argsort(torch.min(labels_index, dim=1)[0], dim=1)
        labels = labels[torch.arange(B).unsqueeze(1), :, sort_index].transpose(-1, -2)

        labels_silence = torch.ones([labels.shape[0], labels.shape[1]]).to(labels) - labels.max(-1)[0]
        labels = torch.cat([labels_silence.unsqueeze(-1), labels], dim=-1)

        labels_nonespk = torch.zeros([B, T, 1]).to(labels)
        labels = torch.cat([labels, labels_nonespk], dim=-1)

        labels = [l[:ilen, :nspk+2] for l, ilen, nspk in zip(labels, clip_lengths, n_spks)]

        # Transform feats to MelSpectrum
        preds, emb_loss, embs, attractors = self.detect(feats, labels, clip_lengths)

        labels_spks = [l[:, 1:-1] for l in labels]
        preds_spks = [p[:, 1:-1] for p in preds]

        max_nspk = max(n_spks)
        labels_spks_pad = pad_labels(labels_spks, max_nspk)
        preds_spks_pad = pad_preds(preds_spks, max_nspk)
        pit_loss1, perm_labels_spks = self.loss_func1(preds_spks_pad, labels_spks_pad, n_spks)
        perm_labels = [torch.cat([l[:, 0].unsqueeze(-1), l_perm], dim=-1) for l, l_perm in zip(labels, perm_labels_spks)]
        perm_labels = [torch.cat([l_perm, l[:, -1].unsqueeze(-1)], dim=-1) for l ,l_perm in zip(labels, perm_labels)]
        pit_loss = self.loss_func2(preds, perm_labels, label_delay=self.label_delay)

        # Calculate the total loss
        tot_loss = pit_loss + emb_loss
        self.log("train/lr", self.opt.param_groups[-1]["lr"], prog_bar=True)
        self.log("train/pit_loss", pit_loss)
        self.log("train/emb_loss", emb_loss)
        self.log("train/tot_loss", tot_loss)

        return tot_loss

    def validation_step(self, batch, batch_index):
        feats, labels, _ = batch
        # Model pred
        clip_lengths = [x.shape[0] for x in feats]
        n_spks = [l.shape[1] for l in labels]
        max_spk = max(n_spks)
        labels = [F.pad(l, (0,max_spk-l.shape[1]), "constant", 0) for l in labels]

        # perm labels
        labels = torch.nn.utils.rnn.pad_sequence(labels, padding_value=0, batch_first=True)
        B, T, max_spk = labels.shape
        frame_index = torch.arange(1, T + 1).to(feats[0]).unsqueeze(0).unsqueeze(-1)
        labels_index = (frame_index * labels)
        labels_index = labels_index.masked_fill_(labels_index == 0, torch.inf)

        # [B, max_spk]
        sort_index = torch.argsort(torch.min(labels_index, dim=1)[0], dim=1)
        labels = labels[torch.arange(B).unsqueeze(1), :, sort_index].transpose(-1, -2)

        labels_silence = torch.ones([labels.shape[0], labels.shape[1]]).to(labels) - labels.max(-1)[0]
        labels = torch.cat([labels_silence.unsqueeze(-1), labels], dim=-1)
        
        labels_nonespk = torch.zeros([B, T, 1]).to(labels)
        labels = torch.cat([labels, labels_nonespk], dim=-1)

        labels = [l[:ilen, :nspk+2] for l, ilen, nspk in zip(labels, clip_lengths, n_spks)]

        preds, emb_loss, embs, attractors = self.detect(feats, labels, clip_lengths)

        labels_spks = [l[:, 1:-1] for l in labels]
        preds_spks = [p[:, 1:-1] for p in preds]

        max_nspk = max(n_spks)
        labels_spks_pad = pad_labels(labels_spks, max_nspk)
        preds_spks_pad = pad_preds(preds_spks, max_nspk)
        pit_loss1, perm_labels_spks = self.loss_func1(preds_spks_pad, labels_spks_pad, n_spks)
        perm_labels = [torch.cat([l[:, 0].unsqueeze(-1), l_perm], dim=-1) for l, l_perm in zip(labels, perm_labels_spks)]
        perm_labels = [torch.cat([l_perm, l[:, -1].unsqueeze(-1)], dim=-1) for l ,l_perm in zip(labels, perm_labels)]
        pit_loss = self.loss_func2(preds, perm_labels, label_delay=self.label_delay)

        # Calculate the total loss
        tot_loss = pit_loss + emb_loss

        # Metrics
        preds_realspk = [p[:, 1:-1] for p in preds]
        labels_realspk = [l[:, 1:-1] for l in perm_labels]
        stats = report_diarization_error(preds_realspk, labels_realspk, label_delay=self.label_delay)

        self.log("val/pit_loss", pit_loss)
        self.log("val/emb_loss", emb_loss)
        self.log("val/tot_loss", tot_loss)

        stats["val_loss"] = [tot_loss.item()]
        return stats
    
    def validation_epoch_end(self, val_step_outputs) -> None:
        stats_holder = defaultdict(list)
        for stats in val_step_outputs:
            for k, v in stats.items():
                stats_holder[k] += v
        # Calculate DER
        stats_avg = {k: sum(v)/len(v) for k, v in stats_holder.items()}
        stats_avg["DER"] = stats_avg["diarization_error"] / stats_avg["speaker_scored"]
        obj_metric = stats_avg['DER']
        self.log("val/obj_metric", obj_metric)
        self.log("val/DER", stats_avg['DER'])
        self.log("val/speech_scored", stats_avg['speech_scored'])
        self.log("val/speech_miss", stats_avg['speech_miss'])
        self.log("val/speech_falarm", stats_avg['speech_falarm'])
        self.log("val/speaker_scored", stats_avg['speaker_scored'])
        self.log("val/speaker_miss", stats_avg['speaker_miss'])
        self.log("val/speaker_falarm", stats_avg['speaker_falarm'])
        self.log("val/speaker_error", stats_avg['speaker_error'])
        self.log("val/diarization_error", stats_avg['diarization_error'])
        self.log("val/frames", stats_avg['frames'])

        # self.scheduler.step(stats_avg["val_loss"])

    def test_step(self, batch, batch_index):
        feats, labels, rec = batch
        # Model pred
        clip_lengths = [x.shape[0] for x in feats]
        n_spks = [l.shape[1] for l in labels]
        max_spk = max(n_spks)
        labels = [F.pad(l, (0,max_spk-l.shape[1]), "constant", 0) for l in labels]

        # perm labels
        labels = torch.nn.utils.rnn.pad_sequence(labels, padding_value=0, batch_first=True)
        B, T, max_spk = labels.shape
        frame_index = torch.arange(1, T + 1).to(feats[0]).unsqueeze(0).unsqueeze(-1)
        labels_index = (frame_index * labels)
        labels_index = labels_index.masked_fill_(labels_index == 0, torch.inf)

        # [B, max_spk]
        sort_index = torch.argsort(torch.min(labels_index, dim=1)[0], dim=1)
        labels = labels[torch.arange(B).unsqueeze(1), :, sort_index].transpose(-1, -2)

        labels_silence = torch.ones([labels.shape[0], labels.shape[1]]).to(labels) - labels.max(-1)[0]
        labels = torch.cat([labels_silence.unsqueeze(-1), labels], dim=-1)

        labels_nonespk = torch.zeros([B, T, 1]).to(labels)
        labels = torch.cat([labels, labels_nonespk], dim=-1)

        labels = [l[:ilen, :nspk+2] for l, ilen, nspk in zip(labels, clip_lengths, n_spks)]
        
        preds, embs, attractors = self.model.test(feats, clip_lengths, self.max_spks + 2)
        
        labels_spks = [l[:, 1:-1] for l in labels]
        preds_spks = [p[:, 1:nspk + 1] for p, nspk in zip(preds, n_spks)]

        max_nspk = max(n_spks)
        labels_spks_pad = pad_labels(labels_spks, max_nspk)
        preds_spks_pad = pad_preds(preds_spks, max_nspk)
        pit_loss1, perm_labels_spks = self.loss_func1(preds_spks_pad, labels_spks_pad, n_spks)
        perm_labels = [torch.cat([l[:, 0].unsqueeze(-1), l_perm], dim=-1) for l, l_perm in zip(labels, perm_labels_spks)]
        perm_labels = [torch.cat([l_perm, l[:, -1].unsqueeze(-1)], dim=-1) for l ,l_perm in zip(labels, perm_labels)]
        # pit_loss = self.loss_func2(preds, perm_labels, label_delay=self.label_delay)

        # Calculate the total loss
        # tot_loss = pit_loss + emb_loss

        # Metrics
        preds_realspk = [p[:, 1:nspk + 1] for p, nspk in zip(preds, n_spks)]
        if preds_realspk[0].shape[1] < max_nspk:
            preds_realspk = pad_preds(preds_realspk, max_nspk)
        labels_realspk = [l[:, 1:-1] for l in perm_labels]
        stats = report_diarization_error(preds_realspk, labels_realspk, label_delay=self.label_delay)

        label_delay = self.hparams["data"]["label_delay"]
        preds = [torch.cat([p[label_delay:, 1:], p[-1, 1:].unsqueeze(0).repeat(label_delay, 1)], dim=0)  for p in preds]
        # num_spks = "2"
        # version = "_ver61_ver_87"
        # save_dir_parnt = f"/data/home/liangdi/dia_projects/pl_version/pl_eend_2.11/tsne_visual/data/onl_{num_spks}spk_version{version}"
        # for content in ["preds", "labels", "embs", "attractors"]:
        #     save_dir = os.path.join(save_dir_parnt, content)
        #     if not os.path.isdir(save_dir):
        #         os.makedirs(save_dir)
        # np.save(f"/data/home/liangdi/dia_projects/pl_version/pl_eend_2.11/tsne_visual/data/onl_{num_spks}spk_version{version}/preds/{rec[0]}.npy", preds[0].detach().cpu().numpy())
        # np.save(f"/data/home/liangdi/dia_projects/pl_version/pl_eend_2.11/tsne_visual/data/onl_{num_spks}spk_version{version}/labels/label_{rec[0]}.npy", labels_realspk[0].detach().cpu().numpy())
        # np.save(f"/data/home/liangdi/dia_projects/pl_version/pl_eend_2.11/tsne_visual/data/onl_{num_spks}spk_version{version}/embs/emb_{rec[0]}.npy", embs[0].detach().cpu().numpy())
        # np.save(f"/data/home/liangdi/dia_projects/pl_version/pl_eend_2.11/tsne_visual/data/onl_{num_spks}spk_version{version}/attractors/attractor_{rec[0]}.npy", attractors[0].detach().cpu().numpy())
        return stats

    def test_epoch_end(self, test_step_outputs) -> None:
        stats_holder = defaultdict(list)
        for stats in test_step_outputs:
            for k, v in stats.items():
                stats_holder[k] += v
        # [print(x) for x in stats_holder["speaker_scored"] if not x ]
        # Calculate DER
        
        stats_avg = {k: sum(v)/len(v) for k, v in stats_holder.items()}
        stats_avg["DER"] = stats_avg["diarization_error"] / stats_avg["speaker_scored"]
        results = {
            "test/DER": stats_avg['DER'],
            "test/speech_scored": stats_avg['speech_scored'],
            "test/speech_miss": stats_avg['speech_miss'],
            "test/speech_falarm": stats_avg['speech_falarm'],
            "test/speaker_scored": stats_avg['speaker_scored'],
            "test/speaker_miss": stats_avg['speaker_miss'],
            "test/speaker_falarm": stats_avg['speaker_falarm'],
            "test/speaker_error": stats_avg['speaker_error'],
            "test/diarization_error": stats_avg['diarization_error'],
            "test/frames": stats_avg['frames'],
            "test/speaker_miss_rate": stats_avg['speaker_miss'] / stats_avg['speaker_scored'],
            "test/speaker_falarm_rate": stats_avg['speaker_falarm'] / stats_avg['speaker_scored'],
            "test/speaker_error_rate": stats_avg['speaker_error'] / stats_avg['speaker_scored'],
        }
        pprint(results)
        self.logger.log_metrics(results)
        self.logger.log_hyperparams(self.hparams, results)
        
        for key in results.keys():
            self.log(key, results[key], prog_bar=True, logger=False)


    def configure_optimizers(self):
        if self.scheduler is not None:
            return {"optimizer": self.opt, "lr_scheduler": {"scheduler": self.scheduler, "monitor": "val/tot_loss", "interval": "epoch"}}
        else:
            return {"optimizer": self.opt}

    def train_dataloader(self):
        return DataLoader(
            self.datasets["train"],
            batch_size=self.hparams["training"]["batch_size"],
            shuffle=self.hparams["training"]["shuffle"],
            num_workers=self.hparams["training"]["n_workers"],
            collate_fn=self.collate_func
        )
    
    def val_dataloader(self):
        return DataLoader(
            self.datasets["val"],
            batch_size=self.hparams["training"]["batch_size"],
            shuffle=False,
            num_workers=self.hparams["training"]["n_workers"],
            collate_fn=self.collate_func
        )
    
    def test_dataloader(self):
        return DataLoader(
            self.datasets["val"],
            batch_size=self.hparams["training"]["batch_size"],
            shuffle=False,
            num_workers=self.hparams["training"]["n_workers"],
            collate_fn=self.collate_func
        )
    
train/oln_tfm_enc_dec_spk_pit.py
Audio-WestlakeU/FS-EEND
# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita)
# Modified by: Yexin Yang
# Modified by: Di Liang
# Licensed under the MIT license.

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
from itertools import permutations
from scipy.optimize import linear_sum_assignment
from torchmetrics import PermutationInvariantTraining
from collections import defaultdict

from torch import Tensor
"""
T: number of frames
C: number of speakers (classes)
D: dimension of embedding (for deep clustering loss)
B: mini-batch size
"""

# labels: [T * C] - T: 50s

def split_results(labels: list[Tensor], preds: list[Tensor], ilens: list[int], blk_size: int):
    n_speakers_split = []
    labels_splt = []
    preds_splt = []
    for l, p, ilen in zip(labels, preds, ilens):
        n_blk = len(l.split(blk_size, dim=0))
        for i in range(n_blk):
            st = i * blk_size
            ed = min((i + 1) * blk_size, ilen)
            non_zero_spk = torch.max(l[:ed, :], dim=0)[0]
            non_zero_spk_idx = non_zero_spk > 0
            nspk = int(non_zero_spk.sum())
            n_speakers_split.append(nspk)
            labels_splt.append(l[st:ed, non_zero_spk_idx])
            preds_splt.append(p[st:ed, :nspk])
    return labels_splt, preds_splt, n_speakers_split

def pad_labels(ts, out_size):
    # pad label's speaker-dim to be model's n_speakers
    for i, t in enumerate(ts):
        if t.shape[1] < out_size:
            # padding
            ts[i] = F.pad(t, (0, out_size - t.shape[1], 0, 0), mode='constant', value=0.)
        elif t.shape[1] > out_size:
            # truncate
            raise ValueError
    return ts


def pad_preds(ys, out_size):
    # pad label's speaker-dim to be model's n_speakers
    ys_padded = []
    for i, y in enumerate(ys):
        if y.shape[1] < out_size:
            # padding
            ys_padded.append(torch.cat([y, torch.zeros((y.shape[0], out_size - y.shape[1]), dtype=y.dtype, device=y.device)], dim=1))
        elif y.shape[1] > out_size:
            # truncate
            raise ValueError
        else:
            ys_padded.append(y)
    return ys_padded

def pit_loss(pred, label, label_delay=0):
    """
    Permutation-invariant training (PIT) cross entropy loss function.

    Args:
      pred:  (T,C)-shaped pre-activation values
      label: (T,C)-shaped labels in {0,1}
      label_delay: if label_delay == 5:
            pred: 0 1 2 3 4 | 5 6 ... 99 100 |
          label: x x x x x | 0 1 ... 94  95 | 96 97 98 99 100
          calculated area: | <------------> |

    Returns:
      min_loss: (1,)-shape mean cross entropy
      label_perms[min_index]: permutated labels
    """
    # label permutations along the speaker axis
    label_perms = [label[..., list(p)] for p
                    in permutations(range(label.shape[-1]))]
    losses = torch.stack(
        [F.binary_cross_entropy_with_logits(
            pred[label_delay:, ...],
            l[:len(l) - label_delay, ...]) for l in label_perms])
    min_loss = losses.min() * (len(label) - label_delay)
    min_index = losses.argmin().detach()
    
    return min_loss, label_perms[min_index]


def batch_pit_loss(ys, ts, label_delay=0):
    """
    PIT loss over mini-batch.

    Args:
      ys: B-length list of predictions
      ts: B-length list of labels

    Returns:
      loss: (1,)-shape mean cross entropy over mini-batch
      labels: B-length list of permuted labels
    """
    loss_w_labels = [pit_loss(y, t, label_delay)
                     for (y, t) in zip(ys, ts)]
    losses, labels = zip(*loss_w_labels)
    loss = torch.stack(losses).sum()
    n_frames = np.sum([t.shape[0] for t in ts])
    loss = loss / n_frames
    return loss, labels


def standard_loss(ys, ts, label_delay=0):
    losses = [F.binary_cross_entropy_with_logits(y[label_delay:, ...], t[:len(t) - label_delay, ...]) * (len(y) - label_delay)
                for (y, t) in zip(ys, ts)]
    loss = torch.stack(losses).sum()
    n_frames = np.sum([t.shape[0] for t in ts]) - (label_delay * len(ts))
    loss = loss / n_frames
    return loss

def standard_mask_loss(ys, ts, label_delay=0):
    losses = []
    for (y, t) in zip(ys, ts):
        loss_mtrx = F.binary_cross_entropy_with_logits(y[label_delay:, ...], t[:len(t) - label_delay, ...], reduce=False)
        loss_mtrx = loss_mtrx.masked_fill((t == 0) & (y < 0), 0)
        loss = loss_mtrx.mean() * (len(y) - label_delay)
        losses.append(loss)
    loss = torch.stack(losses).sum()
    n_frames = np.sum([t.shape[0] for t in ts]) - (label_delay * len(ts))
    loss = loss / n_frames
    return loss


def calc_diarization_error2(pred, label, label_delay=0):
    """
    Calculates diarization error stats for reporting.

    Args:
      pred (torch.FloatTensor): (T,C)-shaped pre-activation values
      label (torch.FloatTensor): (T,C)-shaped labels in {0,1}
      label_delay: if label_delay == 5:
           pred: 0 1 2 3 4 | 5 6 ... 99 100 |
          label: x x x x x | 0 1 ... 94  95 | 96 97 98 99 100
          calculated area: | <------------> |

    Returns:
      res: dict of diarization error stats
    """
    label = label[:len(label) - label_delay, ...]
    decisions = pred[label_delay:, ...] > 0.5
    n_ref = label.sum(axis=-1).long()
    n_sys = decisions.sum(axis=-1).long()
    res = {}
    res['speech_scored'] = (n_ref > 0).sum()
    res['speech_miss'] = ((n_ref > 0) & (n_sys == 0)).sum()
    res['speech_falarm'] = ((n_ref == 0) & (n_sys > 0)).sum()
    res['speaker_scored'] = (n_ref).sum()
    res['speaker_miss'] = torch.max((n_ref - n_sys), torch.zeros_like(n_ref)).sum()
    res['speaker_falarm'] = torch.max((n_sys - n_ref), torch.zeros_like(n_ref)).sum()
    n_map = ((label == 1) & (decisions == 1)).sum(axis=-1)
    res['speaker_error'] = (torch.min(n_ref, n_sys) - n_map).sum()
    res['correct'] = (label == decisions).sum() / label.shape[1]
    res['diarization_error'] = (
        res['speaker_miss'] + res['speaker_falarm'] + res['speaker_error'])
    res['frames'] = len(label)
    # print(res['diarization_error'] / res['speech_scored'], res['speaker_error'] / res['speech_scored'], res['speaker_falarm'] / res['speech_scored'], res['speaker_miss'] / res['speech_scored'])
    # if res['speaker_error'] < 0:
    #     print(decisions)
    #     print(label)
    #     raise Exception("spk error")
    return res



def report_diarization_error2(ys, labels, label_delay=0):
    """
    Reports diarization errors
    Should be called with torch.no_grad

    Args:
      ys: B-length list of predictions (torch.FloatTensor)
      labels: B-length list of labels (torch.FloatTensor)
    """
    stats_batch = defaultdict(list)
    for i, (y, t) in enumerate(zip(ys, labels)):
        stats = calc_diarization_error2(y, t, label_delay)
        for k, v in stats.items():
            stats_batch[k].append(float(v))

    return stats_batch

def calc_diarization_error(pred, label, label_delay=0):
    """
    Calculates diarization error stats for reporting.

    Args:
      pred (torch.FloatTensor): (T,C)-shaped pre-activation values
      label (torch.FloatTensor): (T,C)-shaped labels in {0,1}
      label_delay: if label_delay == 5:
           pred: 0 1 2 3 4 | 5 6 ... 99 100 |
          label: x x x x x | 0 1 ... 94  95 | 96 97 98 99 100
          calculated area: | <------------> |

    Returns:
      res: dict of diarization error stats
    """
    label = label[:len(label) - label_delay, ...]
    decisions = torch.sigmoid(pred[label_delay:, ...]) > 0.5
    n_ref = label.sum(axis=-1).long()
    n_sys = decisions.sum(axis=-1).long()
    res = {}
    res['speech_scored'] = (n_ref > 0).sum()
    res['speech_miss'] = ((n_ref > 0) & (n_sys == 0)).sum()
    res['speech_falarm'] = ((n_ref == 0) & (n_sys > 0)).sum()
    res['speaker_scored'] = (n_ref).sum()
    res['speaker_miss'] = torch.max((n_ref - n_sys), torch.zeros_like(n_ref)).sum()
    res['speaker_falarm'] = torch.max((n_sys - n_ref), torch.zeros_like(n_ref)).sum()
    n_map = ((label == 1) & (decisions == 1)).sum(axis=-1)
    res['speaker_error'] = (torch.min(n_ref, n_sys) - n_map).sum()
    res['correct'] = (label == decisions).sum() / label.shape[1]
    res['diarization_error'] = (
        res['speaker_miss'] + res['speaker_falarm'] + res['speaker_error'])
    res['frames'] = len(label)
    # print(res['diarization_error'] / res['speech_scored'], res['speaker_error'] / res['speech_scored'])
    # print(res['diarization_error'] / res['speaker_scored'])
    if res['speaker_error'] < 0:
        print(decisions)
        print(label)
        raise Exception("spk error")
    return res


def report_diarization_error(ys, labels, label_delay=0):
    """
    Reports diarization errors
    Should be called with torch.no_grad

    Args:
      ys: B-length list of predictions (torch.FloatTensor)
      labels: B-length list of labels (torch.FloatTensor)
    """
    stats_batch = defaultdict(list)
    for i, (y, t) in enumerate(zip(ys, labels)):
        stats = calc_diarization_error(y, t, label_delay)
        for k, v in stats.items():
            stats_batch[k].append(float(v))

    return stats_batch


def batch_pit_n_speaker_loss(ys, ts, n_speakers_list):
    """
    PIT loss over mini-batch.
    Args:
      ys: B-length list of predictions (pre-activations)
      ts: B-length list of labels
      n_speakers_list: list of n_speakers in batch
    Returns:
      loss: (1,)-shape mean cross entropy over mini-batch
      labels: B-length list of permuted labels
    """
    max_n_speakers = max(n_speakers_list)
    # (B, T, C)
    ys = nn.utils.rnn.pad_sequence(ys, padding_value=-1, batch_first=True)

    losses = []
    for shift in range(max_n_speakers):
        # rolled along with speaker-axis
        ts_roll = [torch.roll(t, -shift, dims=1) for t in ts]
        ts_roll = nn.utils.rnn.pad_sequence(ts_roll, padding_value=-1, batch_first=True)
        # loss: (B, T, C)
        loss = F.binary_cross_entropy_with_logits(ys, ts_roll, reduce=False)
        # sum over time: (B, C)
        loss = torch.sum(loss, dim=1)
        losses.append(loss)
    # losses: (B, C, C)
    losses = torch.stack(losses, dim=2)
    # losses[b, i, j] is a loss between
    # `i`-th speaker in y and `(i+j)%C`-th speaker in t
    
    perms = np.array(
        list(permutations(range(max_n_speakers))),
        dtype='i',
    )
    # y_ind: [0,1,2,3]
    y_ind = np.arange(max_n_speakers, dtype='i')
    #  perms  -> relation to t_inds      -> t_inds
    # 0,1,2,3 -> 0+j=0,1+j=1,2+j=2,3+j=3 -> 0,0,0,0
    # 0,1,3,2 -> 0+j=0,1+j=1,2+j=3,3+j=2 -> 0,0,1,3
    t_inds = np.mod(perms - y_ind, max_n_speakers)

    losses_perm = []
    for t_ind in t_inds:
        losses_perm.append(
            torch.mean(losses[:, y_ind, t_ind], dim=1))
    # losses_perm: (B, Perm)
    losses_perm = torch.stack(losses_perm, dim=1)    

    # masks: (B, Perms)
    def select_perm_indices(num, max_num):
        perms = list(permutations(range(max_num)))
        sub_perms = list(permutations(range(num)))
        return [
            [x[:num] for x in perms].index(perm)
            for perm in sub_perms]
    masks = torch.full_like(losses_perm, np.inf)
    for i, t in enumerate(ts):
        n_speakers = n_speakers_list[i]
        indices = select_perm_indices(n_speakers, max_n_speakers)
        masks[i, indices] = 0
    losses_perm += masks

    min_loss = torch.sum(torch.min(losses_perm, dim=1)[0])
    n_frames = np.sum([t.shape[0] for t in ts])
    min_loss = min_loss / n_frames

    min_indices = torch.argmin(losses_perm, dim=1)
    labels_perm = [t[:, perms[idx]] for t, idx in zip(ts, min_indices)]
    labels_perm = [t[:, :n_speakers] for t, n_speakers in zip(labels_perm, n_speakers_list)]

    return min_loss, labels_perm

def batch_pit_n_speaker_loss_label_delay(ys, ts, n_speakers_list, label_delay=0):
    """
    PIT loss over mini-batch.
    Args:
      ys: B-length list of predictions (pre-activations)
      ts: B-length list of labels
      n_speakers_list: list of n_speakers in batch
    Returns:
      loss: (1,)-shape mean cross entropy over mini-batch
      labels: B-length list of permuted labels
    """
    max_n_speakers = max(n_speakers_list)
    ilens = [t.shape[0] for t in ts]
    # (B, T, C)
    ys = nn.utils.rnn.pad_sequence(ys, padding_value=-1, batch_first=True)
    B, T, C = ys.shape


    losses = []
    for shift in range(max_n_speakers):
        # rolled along with speaker-axis
        ts_roll = [torch.roll(t, -shift, dims=1) for t in ts]
        ts_roll = nn.utils.rnn.pad_sequence(ts_roll, padding_value=-1, batch_first=True)
        # loss: (B, T, C)
        loss = F.binary_cross_entropy_with_logits(ys[:, label_delay:, :], ts_roll[:, :T - label_delay, :], reduce=False)
        # sum over time: (B, C)
        for i, ilen in enumerate(ilens):
            loss[i, ilen-label_delay:, :] = 0
        loss = torch.sum(loss, dim=1)
        losses.append(loss)
    # losses: (B, C, C)
    losses = torch.stack(losses, dim=2)
    # losses[b, i, j] is a loss between
    # `i`-th speaker in y and `(i+j)%C`-th speaker in t
    
    perms = np.array(
        list(permutations(range(max_n_speakers))),
        dtype='i',
    )
    # y_ind: [0,1,2,3]
    y_ind = np.arange(max_n_speakers, dtype='i')
    #  perms  -> relation to t_inds      -> t_inds
    # 0,1,2,3 -> 0+j=0,1+j=1,2+j=2,3+j=3 -> 0,0,0,0
    # 0,1,3,2 -> 0+j=0,1+j=1,2+j=3,3+j=2 -> 0,0,1,3
    t_inds = np.mod(perms - y_ind, max_n_speakers)

    losses_perm = []
    for t_ind in t_inds:
        losses_perm.append(
            torch.mean(losses[:, y_ind, t_ind], dim=1))
    # losses_perm: (B, Perm)
    losses_perm = torch.stack(losses_perm, dim=1)    

    # masks: (B, Perms)
    def select_perm_indices(num, max_num):
        perms = list(permutations(range(max_num)))
        sub_perms = list(permutations(range(num)))
        return [
            [x[:num] for x in perms].index(perm)
            for perm in sub_perms]
    masks = torch.full_like(losses_perm, np.inf)
    for i, t in enumerate(ts):
        n_speakers = n_speakers_list[i]
        indices = select_perm_indices(n_speakers, max_n_speakers)
        masks[i, indices] = 0
    losses_perm += masks

    min_loss = torch.sum(torch.min(losses_perm, dim=1)[0])
    n_frames = np.sum([t.shape[0] for t in ts])
    min_loss = min_loss / n_frames

    min_indices = torch.argmin(losses_perm, dim=1)
    labels_perm = [t[:, perms[idx]] for t, idx in zip(ts, min_indices)]
    labels_perm = [t[:, :n_speakers] for t, n_speakers in zip(labels_perm, n_speakers_list)]

    return min_loss, labels_perm

def batch_pit_n_speaker_loss2(ys, ts, n_speakers_list):
    """
    PIT loss over mini-batch.
    Args:
      ys: B-length list of predictions (pre-activations)
      ts: B-length list of labels
      n_speakers_list: list of n_speakers in batch
    Returns:
      loss: (1,)-shape mean cross entropy over mini-batch
      labels: B-length list of permuted labels
    """
    max_n_speakers = max(n_speakers_list)
    # (B, T, C)
    ys = nn.utils.rnn.pad_sequence(ys, padding_value=-1, batch_first=True)

    losses = []
    for shift in range(max_n_speakers):
        # rolled along with speaker-axis
        ts_roll = [torch.roll(t, -shift, dims=1) for t in ts]
        ts_roll = nn.utils.rnn.pad_sequence(ts_roll, padding_value=-1, batch_first=True)
        # loss: (B, T, C)
        loss = F.binary_cross_entropy(ys, ts_roll, reduce=False)
        # sum over time: (B, C)
        loss = torch.sum(loss, dim=1)
        losses.append(loss)
    # losses: (B, C, C)
    losses = torch.stack(losses, dim=2)
    # losses[b, i, j] is a loss between
    # `i`-th speaker in y and `(i+j)%C`-th speaker in t
    
    perms = np.array(
        list(permutations(range(max_n_speakers))),
        dtype='i',
    )
    # y_ind: [0,1,2,3]
    y_ind = np.arange(max_n_speakers, dtype='i')
    #  perms  -> relation to t_inds      -> t_inds
    # 0,1,2,3 -> 0+j=0,1+j=1,2+j=2,3+j=3 -> 0,0,0,0
    # 0,1,3,2 -> 0+j=0,1+j=1,2+j=3,3+j=2 -> 0,0,1,3
    t_inds = np.mod(perms - y_ind, max_n_speakers)

    losses_perm = []
    for t_ind in t_inds:
        losses_perm.append(
            torch.mean(losses[:, y_ind, t_ind], dim=1))
    # losses_perm: (B, Perm)
    losses_perm = torch.stack(losses_perm, dim=1)    

    # masks: (B, Perms)
    def select_perm_indices(num, max_num):
        perms = list(permutations(range(max_num)))
        sub_perms = list(permutations(range(num)))
        return [
            [x[:num] for x in perms].index(perm)
            for perm in sub_perms]
    masks = torch.full_like(losses_perm, np.inf)
    for i, t in enumerate(ts):
        n_speakers = n_speakers_list[i]
        indices = select_perm_indices(n_speakers, max_n_speakers)
        masks[i, indices] = 0
    losses_perm += masks

    min_loss = torch.sum(torch.min(losses_perm, dim=1)[0])
    n_frames = np.sum([t.shape[0] for t in ts])
    min_loss = min_loss / n_frames

    min_indices = torch.argmin(losses_perm, dim=1)
    labels_perm = [t[:, perms[idx]] for t, idx in zip(ts, min_indices)]
    labels_perm = [t[:, :n_speakers] for t, n_speakers in zip(labels_perm, n_speakers_list)]

    return min_loss, labels_perm

class PITLoss(nn.Module):
    def __init__(self, n_spks=4) -> None:
        super(PITLoss, self).__init__()
        self.n_spks = n_spks
        trav_idx = []
        for x in permutations(list(range(n_spks))):
            trav_idx += x
        self.trav_idx = torch.tensor(trav_idx).long()
    
    def forward(self, preds, labels):
        B, T, _ = preds.shape

        labels_all_case = labels[:, :, self.trav_idx].reshape(B, T, -1, self.n_spks)
        case_num = labels_all_case.shape[-2]
        preds_all_case = preds.unsqueeze(-1).repeat(1, 1, 1, case_num).transpose(-1, -2)

        loss_all_case = F.binary_cross_entropy_with_logits(preds_all_case, labels_all_case, reduction="none").mean(-1).mean(1)

        min_idx = torch.argmin(loss_all_case, dim=-1)
        selected_loss = loss_all_case[torch.arange(B), min_idx]
        pit_min_loss = selected_loss.sum() / B

        perm_labels = labels_all_case[torch.arange(B), :, min_idx, :]
        return pit_min_loss, perm_labels

    def swap(self, a, i, j):
        temp = a[i]
        a[i] = a[j]
        a[j] = temp

    def dfs(self, a, depth=0):
        r = []
        if depth == len(a):
            r += a
        for i in range(depth, len(a)):
            self.swap(a, i, depth)
            r += self.dfs(a, depth + 1)
            self.swap(a, depth, i)
        return r

if __name__ == "__main__":
    torch.random.manual_seed(2)
    pred = torch.rand([5, 1000, 4]).float()
    label = torch.randint(0, 2, [5, 1000, 4]).float()
    x_1, y_1 = batch_pit_loss(pred, label)

    x_2, y_2 = batch_pit_n_speaker_loss(pred, label)
    print("X:", x_1.item() == x_2.item())
    y_1_t = torch.stack(y_1)
    y_2_t = torch.stack(y_2)
    print("Y:", (y_1_t == y_2_t).float().mean())

    print(report_diarization_error(y_1, label))
    print(report_diarization_error(y_2, label))


train/utils/loss.py
Audio-WestlakeU/FS-EEND
import tqdm
import torch
from torch import Tensor
import numpy as np
import torch.nn.functional as F
from numpy import random as nr

from scipy.optimize import linear_sum_assignment

from torch.utils.data import WeightedRandomSampler

def split_inp(feats: list[Tensor], labels: list[Tensor], T_prime: int):
    # fea: (T, D)
    rechunk_feats = []
    rechunk_labels = []
    for fea, l in zip(feats, labels):
        rechk_fea = list(fea.split(T_prime, dim=0))
        rechunk_feats = rechunk_feats + rechk_fea
        rechk_l = list(l.split(T_prime, dim=0))
        rechunk_labels = rechunk_labels + rechk_l
    return rechunk_feats, rechunk_labels


def resize_chunk(T: int)->int:
    var_chunks = np.array([50, 100, 200, 500, 1000])
    T_prime = T
    if torch.rand(1) >= 0.5:
        T_prime = min(nr.choice(var_chunks, size=(1), replace=False).tolist()[0], T)
    return T_prime


def upd_buf_ver2(x_buf, x_i, z_buf, y_i, buf_size):
    # (T, S_i)
    x_cat = torch.cat([x_buf, x_i], dim=0)
    y_cat = torch.cat([z_buf, y_i], dim=0)

    T, S_i = y_cat.shape
    p = y_cat / y_cat.sum(dim=1, keepdim=True)
    p = p.masked_fill_(p == 0, 1e-6)
    q = torch.tensor(1.0 / S_i).repeat(T, S_i).to(p)
    r = torch.sum(y_cat / y_cat.sum(dim=0, keepdim=True), dim=1)
    KLD = r * ((p * torch.log(p / q)).sum(dim=1))

    prob = KLD / KLD.sum(dim=0, keepdim=True)
    # prob_cumsum = F.pad(prob.cumsum(dim=0), (1, 0), mode='constant', value=0.)

    # t_selct = []
    # for i in range(buf_size):
    #     rand_prob = torch.rand(1).to(prob_cumsum)
    #     for t in range(1, T + 1):
    #         if rand_prob >= prob_cumsum[t-1] and rand_prob <= prob_cumsum[t]:
    #             t_selct.append(t - 1)
    t_selct = list(WeightedRandomSampler(prob, buf_size, replacement=False))
    t_selct.sort()
    x_buf_upd = x_cat[t_selct]
    y_buf_upd = y_cat[t_selct]

    return x_buf_upd, y_buf_upd

def upd_buf(x_buf, x_i, z_buf, y_i, buf_size):
    # (T, S_i)
    x_cat = torch.cat([x_buf, x_i], dim=0)
    y_cat = torch.cat([z_buf, y_i], dim=0)

    T, S_i = y_cat.shape
    p = y_cat / y_cat.sum(dim=1, keepdim=True)
    p = p.masked_fill_(p == 0, 1e-6)
    q = torch.tensor(1.0 / S_i).repeat(T, S_i).to(p)
    KLD =  (p * torch.log(p / q)).sum(dim=1)
    KLD = KLD.masked_fill_(KLD<0, 0)
    KLD = KLD.masked_fill_(KLD==0, 1e-6)

    prob = KLD / KLD.sum(dim=0, keepdim=True)
    # prob_cumsum = F.pad(prob.cumsum(dim=0), (1, 0), mode='constant', value=0.)

    # t_selct = []
    # for i in range(buf_size):
    #     rand_prob = torch.rand(1).to(prob_cumsum)
    #     for t in range(1, T + 1):
    #         if rand_prob >= prob_cumsum[t-1] and rand_prob <= prob_cumsum[t]:
    #             t_selct.append(t - 1)
    t_selct = list(WeightedRandomSampler(prob, buf_size, replacement=False))
    t_selct.sort()
    x_buf_upd = x_cat[t_selct]
    y_buf_upd = y_cat[t_selct]

    return x_buf_upd, y_buf_upd

def cal_cor(A: Tensor, B: Tensor):
    # (T)
    A_mean = A.mean()
    B_mean = B.mean()
    cov = torch.sum((A - A_mean) * (B - B_mean))
    stdv1 = torch.sqrt(torch.sum((A - A_mean) * (A - A_mean)))
    stdv2 = torch.sqrt(torch.sum((B - B_mean) * (B - B_mean)))
    return cov / (stdv1 * stdv2 + 1e-6)



def find_best_perm(y: Tensor, y_pred: Tensor):
    # y(T, C) # y_pred (T, C)
    T, C = y.shape
    cc_mtrx = torch.zeros(C, C)
    row = torch.arange(C)
    for i in range(y.shape[1]):
        col = row.roll(shifts=i, dims=0)
        for ro, co in zip(row, col):
            cc_mtrx[ro, co] = cal_cor(y[:, ro], y_pred[:, co])
    
    best_perm = linear_sum_assignment(cc_mtrx, True)[1]
    return best_perm


class ContextBuilder(torch.nn.Module):
    def __init__(self, context_size) -> None:
        super(ContextBuilder, self).__init__()
        self.context_size = context_size

    def forward(self, x):
        bsz, T, _ = x.shape
        x_pad = torch.nn.functional.pad(x, (0, 0, self.context_size, self.context_size))
        return x_pad.unfold(1, T, 1).reshape(bsz, -1, T).transpose(-1, -2)


class TorchScaler(torch.nn.Module):
    """
    This torch module implements scaling for input tensors, both instance based
    and dataset-wide statistic based.

    Args:
        statistic: str, (default='dataset'), represent how to compute the statistic for normalisation.
            Choice in {'dataset', 'instance'}.
             'dataset' needs to be 'fit()' with a dataloader of the dataset.
             'instance' apply the normalisation at an instance-level, so compute the statitics on the instance
             specified, it can be a clip or a batch.
        normtype: str, (default='standard') the type of normalisation to use.
            Choice in {'standard', 'mean', 'minmax'}. 'standard' applies a classic normalisation with mean and standard
            deviation. 'mean' substract the mean to the data. 'minmax' substract the minimum of the data and divide by
            the difference between max and min.
    """

    def __init__(self, statistic="dataset", normtype="standard", dims=(1, 2), eps=1e-8):
        super(TorchScaler, self).__init__()
        assert statistic in ["dataset", "instance"]
        assert normtype in ["standard", "mean", "minmax"]
        if statistic == "dataset" and normtype == "minmax":
            raise NotImplementedError(
                "statistic==dataset and normtype==minmax is not currently implemented."
            )
        self.statistic = statistic
        self.normtype = normtype
        self.dims = dims
        self.eps = eps

    def load_state_dict(self, state_dict, strict=True):
        if self.statistic == "dataset":
            super(TorchScaler, self).load_state_dict(state_dict, strict)

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        if self.statistic == "dataset":
            super(TorchScaler, self)._load_from_state_dict(
                state_dict,
                prefix,
                local_metadata,
                strict,
                missing_keys,
                unexpected_keys,
                error_msgs,
            )

    def fit(self, dataloader, transform_func=lambda x: x[0]):
        """
        Scaler fitting

        Args:
            dataloader (DataLoader): training data DataLoader
            transform_func (lambda function, optional): Transforms applied to the data.
                Defaults to lambdax:x[0].
        """
        indx = 0
        for batch in tqdm.tqdm(dataloader):

            feats = transform_func(batch)
            if indx == 0:
                mean = torch.mean(feats, self.dims, keepdim=True).mean(0).unsqueeze(0)
                mean_squared = (
                    torch.mean(feats ** 2, self.dims, keepdim=True).mean(0).unsqueeze(0)
                )
            else:
                mean += torch.mean(feats, self.dims, keepdim=True).mean(0).unsqueeze(0)
                mean_squared += (
                    torch.mean(feats ** 2, self.dims, keepdim=True).mean(0).unsqueeze(0)
                )
            indx += 1

        mean /= indx
        mean_squared /= indx

        self.register_buffer("mean", mean)
        self.register_buffer("mean_squared", mean_squared)

    def forward(self, tensor):
        if self.statistic == "dataset":
            assert hasattr(self, "mean") and hasattr(
                self, "mean_squared"
            ), "TorchScaler should be fit before used if statistics=dataset"
            assert tensor.ndim == self.mean.ndim, "Pre-computed statistics "
            if self.normtype == "mean":
                return tensor - self.mean
            elif self.normtype == "standard":
                std = torch.sqrt(self.mean_squared - self.mean ** 2)
                return (tensor - self.mean) / (std + self.eps)
            else:
                raise NotImplementedError

        else:
            if self.normtype == "mean":
                return tensor - torch.mean(tensor, self.dims, keepdim=True)
            elif self.normtype == "standard":
                return (tensor - torch.mean(tensor, self.dims, keepdim=True)) / (
                    torch.std(tensor, self.dims, keepdim=True) + self.eps
                )
            elif self.normtype == "minmax":
                return (tensor - torch.amin(tensor, dim=self.dims, keepdim=True)) / (
                    torch.amax(tensor, dim=self.dims, keepdim=True)
                    - torch.amin(tensor, dim=self.dims, keepdim=True)
                    + self.eps
                )

def splice(Y, context_size=0):
    """ Frame splicing

    Args:
        Y: feature
            (n_frames, n_featdim)-shaped numpy array
        context_size:
            number of frames concatenated on left-side
            if context_size = 5, 11 frames are concatenated.

    Returns:
        Y_spliced: spliced feature
            (n_frames, n_featdim * (2 * context_size + 1))-shaped
    """
    Y_pad = np.pad(
            Y,
            [(context_size, context_size), (0, 0)],
            'constant')
    Y_spliced = np.lib.stride_tricks.as_strided(
            Y_pad,
            (Y.shape[0], Y.shape[1] * (2 * context_size + 1)),
            (Y.itemsize * Y.shape[1], Y.itemsize), writeable=False)
    return Y_spliced

# if __name__ == "__main__":
#     import numpy as np

#     context = ContextBuilder(7)
#     data = torch.rand([1, 10, 2])
#     data_contx = context(data)

#     data_np = data.numpy()[0]
#     data_contx_np = splice(data_np, 7)

#     print(data_contx.numpy()[0] == data_contx_np)

if __name__=="__main__":
    a = torch.tensor([[1.0,2.0], [3.0,4.0], [2.0,1.0]])
    b = torch.tensor([[1.0,1.0], [2.0,4.0], [2.0,2.0]])
    find_best_perm(a, b)
train/utils/utils.py
Audio-WestlakeU/FS-EEND
import os
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument('--OMP_NUM_THREADS', type=int, default=1)
temp_args, _ = parser.parse_known_args()
os.environ["OMP_NUM_THREADS"] = str(temp_args.OMP_NUM_THREADS)

import random
import torch
import yaml
import hyperpyyaml

import numpy as np
import pytorch_lightning as pl

from pytorch_lightning.profilers import AdvancedProfiler

from functools import partial
from collections import defaultdict
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger
from nnet.model.onl_tfm_enc_1dcnn_enc_linear_non_autoreg_pos_enc_l2norm import OnlineTransformerDADiarization
from utlis.scheduler import NoamScheduler
from datasets.diarization_dataset import KaldiDiarizationDataset, my_collate
from train.oln_tfm_enc_dec import SpeakerDiarization

import warnings
warnings.filterwarnings("ignore")


def train(configs, gpus, checkpoint_resume, test_folder=None):

    train_set = KaldiDiarizationDataset(
            data_dir=configs["data"]["train_data_dir"],
            chunk_size=configs["data"]["chunk_size"],
            context_size=configs["data"]["context_recp"],
            input_transform=configs["data"]["feat_type"],
            frame_size=configs["data"]["feat"]["win_length"],
            frame_shift=configs["data"]["feat"]["hop_length"],
            subsampling=configs["data"]["subsampling"],
            rate=configs["data"]["feat"]["sample_rate"],
            label_delay=configs["data"]["label_delay"],
            n_speakers=configs["data"]["num_speakers"],
            use_last_samples=configs["data"]["use_last_samples"],
            shuffle=configs["data"]["shuffle"])

    val_set = KaldiDiarizationDataset(
            data_dir=configs["data"]["val_data_dir"],
            chunk_size=configs["data"]["chunk_size"],
            context_size=configs["data"]["context_recp"],
            input_transform=configs["data"]["feat_type"],
            frame_size=configs["data"]["feat"]["win_length"],
            frame_shift=configs["data"]["feat"]["hop_length"],
            subsampling=configs["data"]["subsampling"],
            rate=configs["data"]["feat"]["sample_rate"],
            label_delay=configs["data"]["label_delay"],
            n_speakers=configs["data"]["num_speakers"],
            use_last_samples=configs["data"]["use_last_samples"],
            shuffle=configs["data"]["shuffle"])
    
    datasets = {
        "train": train_set,
        "val": val_set
    }
    
    collate_func = my_collate

    # Define model
    model = OnlineTransformerDADiarization(
        n_speakers=configs["data"]["num_speakers"],
        in_size=(2 * configs["data"]["context_recp"] + 1) * configs["data"]["feat"]["n_mels"],          # Transformer need to know maximum data length
        **configs["model"]["params"],
    )


    # Define optimizer
    opt_config = {
        "params": model.parameters(),
        "lr": configs["training"]["lr"]
    }
    opt_name = configs["training"]["opt"].lower()
    if opt_name == "adam":
        opt = torch.optim.Adam
    elif opt_name == "sgd":
        opt = torch.optim.SGD
    elif opt_name == "noam":
        opt = partial(
            torch.optim.Adam,
            betas=(0.9, 0.98),
            eps=1e-9
        )
    else: 
        NotImplementedError
    opt = opt(**opt_config)

    if configs["training"]["scheduler"]:
        print("Using noam scheduler")
        scheduler = NoamScheduler(opt, configs["model"]["params"]["n_units"], configs["training"]["warm_steps"], scale=configs["training"]["schedule_scale"]) if configs["training"]["scheduler"].lower() == "noam" else NotImplementedError
    else:
        scheduler = None

    # Define the logger
    logger = TensorBoardLogger(os.path.dirname(configs["log"]["log_dir"]), configs["log"]["model_name"])
    configs["log"]["log_dir"] = logger.log_dir  # updarte log_dir to: ./logs/model_xx/version_xx
    print("Experiment dir:", configs["log"]["log_dir"])
    os.makedirs(configs["log"]["log_dir"], exist_ok=True)
    with open(configs["log"]["log_dir"] + "/config.yaml", "w") as f:
        docs = yaml.dump(configs, f)
        f.close()
    callbacks = [
        EarlyStopping(monitor="val/obj_metric", patience=configs["training"]["early_stop_epoch"], verbose=True, mode="min"),
        ModelCheckpoint(logger.log_dir, monitor="val/obj_metric", save_top_k=configs["log"]["save_top_k"], mode="min", save_last=True)
    ]


    # Define the training setup
    spk_dia_main = SpeakerDiarization(
        hparams=configs,
        model=model,
        datasets=datasets,
        opt=opt,
        scheduler=scheduler,
        collate_func=collate_func
    )

    # Initialization
    if configs["training"]["init_ckpt"]:
        print("Load from checkpoint {} ... ".format(configs["training"]["init_ckpt"]))
        ckpt_package = torch.load(configs["training"]["init_ckpt"], map_location="cpu")
        # state_dict = ckpt_package["state_dict"]
        spk_dia_main.load_state_dict(ckpt_package)
    
    # Freezing parameters
    # Freeze all
    # for name, params in model.named_parameters():
    #     if "dec" in name:
    #         print("Unfreeze:", name)
    #         params.requires_grad = True
    #     else:
    #         print("Freeze:", name)
    #         params.requires_grad = False

    # Define the trainer
    profiler = AdvancedProfiler(filename="perf_logs")
    trainer = pl.Trainer(
        max_epochs=configs["training"]["max_epochs"],
        callbacks=callbacks,
        gpus=gpus,
        strategy=configs["training"]["dist_strategy"],
        accumulate_grad_batches=configs["training"]["grad_accm"],
        logger=logger,
        resume_from_checkpoint=checkpoint_resume,
        gradient_clip_val=configs["training"]["grad_clip"],
        check_val_every_n_epoch=configs["training"]["val_interval"],
        **configs["debug"]
    )

    if test_folder is None:
        # Start training
        trainer.fit(spk_dia_main)
        best_path = trainer.checkpoint_callback.best_model_path
        print("Best model path:", best_path)
        test_folder = os.path.dirname(best_path)
    
    # Test step: given the folder and average the models in the given folder
    for _, _, files in os.walk(test_folder):
        all_files = files
    # ckpts = [x for x in all_files if (".ckpt" in x) and ("epoch" in x)]
    ckpts = [x for x in all_files if (".ckpt" in x) and ("epoch" in x) and int(x.split("=")[1].split("-")[0])>=configs["log"]["start_epoch"] and int(x.split("=")[1].split("-")[0])<=configs["log"]["end_epoch"]]

    # [os.remove(test_folder + "/" + x) for x in all_files if (".ckpt" in x) and ("epoch" in x) and (int(x.split("=")[1].split("-")[0])<=39 or int(x.split("=")[1].split("-")[0])>=50 and int(x.split("=")[1].split("-")[0])<=89)]

    print("Test using ckpts:")
    [print(test_folder + "/" + x) for x in ckpts]
    test_state = defaultdict(float)
    for c in ckpts:
        state_dict = torch.load(test_folder + "/" + c, map_location=torch.device("cuda:{}".format(gpus[0])))["state_dict"]
        for name, param in state_dict.items():
            test_state[name] += param / len(ckpts)

    if configs["log"]["save_avg_path"]:
        torch.save(test_state, configs["log"]["save_avg_path"])
    spk_dia_main.load_state_dict(test_state)
    trainer.test(spk_dia_main)

if __name__ == "__main__":

    parser = ArgumentParser()

    parser.add_argument('--configs', help='Configuration file path', required=True)
    parser.add_argument('--gpus', default=None, help='Device used for training')
    parser.add_argument("--checkpoint_resume", default=None, help="Checkpoint path to resume training")
    parser.add_argument("--test_from_folder", default=None, help="Checkpoint path to test training")
    setup = parser.parse_args()
    with open(setup.configs, "r") as f:
        configs = hyperpyyaml.load_hyperpyyaml(f)
        f.close()

    # Freeze seed
    seed = configs["training"]["seed"]
    if seed:
        torch.random.manual_seed(seed)
        np.random.seed(seed)
        random.seed(seed)
        pl.seed_everything(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True

    train(configs, gpus=setup.gpus, checkpoint_resume=setup.checkpoint_resume, test_folder=setup.test_from_folder)


    

train_dia.py
Audio-WestlakeU/FS-EEND
import os
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument('--OMP_NUM_THREADS', type=int, default=1)
temp_args, _ = parser.parse_known_args()
os.environ["OMP_NUM_THREADS"] = str(temp_args.OMP_NUM_THREADS)

import random
import torch
import yaml
import hyperpyyaml

import numpy as np
import pytorch_lightning as pl

from pytorch_lightning.profilers import AdvancedProfiler

from functools import partial
from collections import defaultdict
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger
from nnet.model.onl_tfm_enc_1dcnn_enc_linear_non_autoreg_pos_enc_l2norm import OnlineTransformerDADiarization
from utlis.scheduler import NoamScheduler
from datasets.dummy_dataset import KaldiDiarizationDataset, my_collate
from train.oln_tfm_enc_dec_spk_pit import SpeakerDiarization

import warnings
warnings.filterwarnings("ignore")


def train(configs, gpus, checkpoint_resume, test_folder=None):

    train_set = KaldiDiarizationDataset(
            data_dir=configs["data"]["train_data_dir"],
            chunk_size=configs["data"]["chunk_size"],
            context_size=configs["data"]["context_recp"],
            input_transform=configs["data"]["feat_type"],
            frame_size=configs["data"]["feat"]["win_length"],
            frame_shift=configs["data"]["feat"]["hop_length"],
            subsampling=configs["data"]["subsampling"],
            rate=configs["data"]["feat"]["sample_rate"],
            label_delay=configs["data"]["label_delay"],
            n_speakers=configs["data"]["num_speakers"],
            use_last_samples=configs["data"]["use_last_samples"],
            shuffle=configs["data"]["shuffle"])

    val_set = KaldiDiarizationDataset(
            data_dir=configs["data"]["val_data_dir"],
            chunk_size=configs["data"]["chunk_size"],
            context_size=configs["data"]["context_recp"],
            input_transform=configs["data"]["feat_type"],
            frame_size=configs["data"]["feat"]["win_length"],
            frame_shift=configs["data"]["feat"]["hop_length"],
            subsampling=configs["data"]["subsampling"],
            rate=configs["data"]["feat"]["sample_rate"],
            label_delay=configs["data"]["label_delay"],
            n_speakers=configs["data"]["num_speakers"],
            use_last_samples=configs["data"]["use_last_samples"],
            shuffle=configs["data"]["shuffle"])
    
    datasets = {
        "train": train_set,
        "val": val_set
    }
    
    collate_func = my_collate

    # Define model
    model = OnlineTransformerDADiarization(
        n_speakers=configs["data"]["num_speakers"],
        in_size=(2 * configs["data"]["context_recp"] + 1) * configs["data"]["feat"]["n_mels"],          # Transformer need to know maximum data length
        **configs["model"]["params"],
    )


    # Define optimizer
    opt_config = {
        "params": model.parameters(),
        "lr": configs["training"]["lr"]
    }
    opt_name = configs["training"]["opt"].lower()
    if opt_name == "adam":
        opt = torch.optim.Adam
    elif opt_name == "sgd":
        opt = torch.optim.SGD
    elif opt_name == "noam":
        opt = partial(
            torch.optim.Adam,
            betas=(0.9, 0.98),
            eps=1e-9
        )
    else: 
        NotImplementedError
    opt = opt(**opt_config)

    if configs["training"]["scheduler"]:
        print("Using noam scheduler")
        scheduler = NoamScheduler(opt, configs["model"]["params"]["n_units"], configs["training"]["warm_steps"], scale=configs["training"]["schedule_scale"]) if configs["training"]["scheduler"].lower() == "noam" else NotImplementedError
    else:
        scheduler = None

    # Define the logger
    logger = TensorBoardLogger(os.path.dirname(configs["log"]["log_dir"]), configs["log"]["model_name"])
    configs["log"]["log_dir"] = logger.log_dir  # updarte log_dir to: ./logs/model_xx/version_xx
    print("Experiment dir:", configs["log"]["log_dir"])
    os.makedirs(configs["log"]["log_dir"], exist_ok=True)
    with open(configs["log"]["log_dir"] + "/config.yaml", "w") as f:
        docs = yaml.dump(configs, f)
        f.close()
    callbacks = [
        EarlyStopping(monitor="val/obj_metric", patience=configs["training"]["early_stop_epoch"], verbose=True, mode="min"),
        ModelCheckpoint(logger.log_dir, monitor="val/obj_metric", save_top_k=configs["log"]["save_top_k"], mode="min", save_last=True)
    ]


    # Define the training setup
    spk_dia_main = SpeakerDiarization(
        hparams=configs,
        model=model,
        datasets=datasets,
        opt=opt,
        scheduler=scheduler,
        collate_func=collate_func
    )

    # Initialization
    if configs["training"]["init_ckpt"]:
        print("Load from checkpoint {} ... ".format(configs["training"]["init_ckpt"]))
        ckpt_package = torch.load(configs["training"]["init_ckpt"], map_location="cpu")
        # state_dict = ckpt_package["state_dict"]
        spk_dia_main.load_state_dict(ckpt_package)
    
    # Freezing parameters
    # Freeze all
    # for name, params in model.named_parameters():
    #     if "dec" in name:
    #         print("Unfreeze:", name)
    #         params.requires_grad = True
    #     else:
    #         print("Freeze:", name)
    #         params.requires_grad = False

    # Define the trainer
    profiler = AdvancedProfiler(filename="perf_logs")
    trainer = pl.Trainer(
        max_epochs=configs["training"]["max_epochs"],
        callbacks=callbacks,
        gpus=gpus,
        strategy=configs["training"]["dist_strategy"],
        accumulate_grad_batches=configs["training"]["grad_accm"],
        logger=logger,
        resume_from_checkpoint=checkpoint_resume,
        gradient_clip_val=configs["training"]["grad_clip"],
        check_val_every_n_epoch=configs["training"]["val_interval"],
        **configs["debug"]
    )

    if test_folder is None:
        # Start training
        trainer.fit(spk_dia_main)
        best_path = trainer.checkpoint_callback.best_model_path
        print("Best model path:", best_path)
        test_folder = os.path.dirname(best_path)
    
    # Test step: given the folder and average the models in the given folder
    for _, _, files in os.walk(test_folder):
        all_files = files
    # ckpts = [x for x in all_files if (".ckpt" in x) and ("epoch" in x)]
    ckpts = [x for x in all_files if (".ckpt" in x) and ("epoch" in x) and int(x.split("=")[1].split("-")[0])>=configs["log"]["start_epoch"] and int(x.split("=")[1].split("-")[0])<=configs["log"]["end_epoch"]]

    # [os.remove(test_folder + "/" + x) for x in all_files if (".ckpt" in x) and ("epoch" in x) and (int(x.split("=")[1].split("-")[0])<=39 or int(x.split("=")[1].split("-")[0])>=50 and int(x.split("=")[1].split("-")[0])<=89)]

    print("Test using ckpts:")
    [print(test_folder + "/" + x) for x in ckpts]
    test_state = defaultdict(float)
    for c in ckpts:
        # state_dict = torch.load(test_folder + "/" + c, map_location=torch.device("cuda:{}".format(gpus[0])))["state_dict"]
        state_dict = torch.load(test_folder + "/" + c, map_location="cpu")["state_dict"]
        for name, param in state_dict.items():
            test_state[name] += param / len(ckpts)

    if configs["log"]["save_avg_path"]:
        torch.save(test_state, configs["log"]["save_avg_path"])
    spk_dia_main.load_state_dict(test_state)
    trainer.test(spk_dia_main)

if __name__ == "__main__":

    parser = ArgumentParser()

    parser.add_argument('--configs', help='Configuration file path', required=True)
    parser.add_argument('--gpus', default=None, help='Device used for training')
    parser.add_argument("--checkpoint_resume", default=None, help="Checkpoint path to resume training")
    parser.add_argument("--test_from_folder", default=None, help="Checkpoint path to test training")
    setup = parser.parse_args()
    with open(setup.configs, "r") as f:
        configs = hyperpyyaml.load_hyperpyyaml(f)
        f.close()

    # Freeze seed
    seed = configs["training"]["seed"]
    if seed:
        torch.random.manual_seed(seed)
        np.random.seed(seed)
        random.seed(seed)
        pl.seed_everything(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True

    train(configs, gpus=setup.gpus, checkpoint_resume=setup.checkpoint_resume, test_folder=setup.test_from_folder)


    

train_dia_fintn.py
Audio-WestlakeU/FS-EEND
from torch.optim.lr_scheduler import _LRScheduler

class NoamScheduler(_LRScheduler):
    """
    See https://arxiv.org/pdf/1706.03762.pdf
    lrate = d_model**(-0.5) * \
            min(step_num**(-0.5), step_num*warmup_steps**(-1.5))
    Args:
        d_model: int
            The number of expected features in the encoder inputs.
        warmup_steps: int
            The number of steps to linearly increase the learning rate.
    """
    def __init__(self, optimizer, d_model, warmup_steps, scale, last_epoch=-1):
        self.d_model = d_model
        self.warmup_steps = warmup_steps
        self.scale = scale
        super(NoamScheduler, self).__init__(optimizer, last_epoch)
        print("Noam initializing...")
        if self.last_epoch == -1:
            for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
                param_group['lr'] = lr
            self.last_epoch = 0
    
    def get_lr(self):
        last_epoch = max(1, self.last_epoch)
        scale = self.scale * self.d_model ** (-0.5) * min(last_epoch ** (-0.5), last_epoch * self.warmup_steps ** (-1.5))
        return [base_lr * scale for base_lr in self.base_lrs]
utlis/scheduler.py
Audio-WestlakeU/FS-EEND
import os
import numpy as np
import h5py
import torch

out_root = "/data/home/liangdi/dia_projects/pl_version/pl_eend_2.11/tsne_visual/data/onl_4spk_version_ver37_ver_73/preds"
out_h5_root = "/data/home/liangdi/dia_projects/pl_version/pl_eend_2.11/tsne_visual/data/onl_4spk_version_ver37_ver_73/preds_h5"
if not os.path.isdir(out_h5_root):
    os.mkdir(out_h5_root)
bsz = 500

for dir_path, dir_names, file_names in os.walk(out_root):
    for file_name in file_names:
        pred = np.load(os.path.join(out_root, file_name))
        pred = 1 / (1 + np.exp(-pred))
        print(pred)
        name = file_name.split('.')[0]
        h5_file = h5py.File(out_h5_root + f"/{name}.h5", "w")
        h5_file.create_dataset("T_hat", data=pred)
        h5_file.close()

visualize/gen_h5_output.py
Audio-WestlakeU/FS-EEND