[ASR] Add optimization util for linear sum assignment algorithm #6349

tango4j · 2023-04-03T02:07:07Z

What does this PR do ?

Add a optimization util file for linear sum assignment (LSA) solver algorithm for online-diarization/multi-speaker-ASR
LSA problem solver is needed for the following tasks in NeMo:
(1) Permutation Invariant Loss (PIL) for diarization model training
(2) Label permutation matching for online speaker diarzation
(3) Concatenated minimum-permutation Word Error Rate (cp-WER) calculation

What is LSA solver algorithm? Google OR-tools LSA Solver

The NeMo linear_sum_assignment function is compared with scipy.optimization.linear_sum_assingment.
In the unit-test for NeMo LSA solver, the result is compared with the scipy version of linear_sum_assignment.

Removing @torch.jit.script decorator in speaker_utils.py since it creates type-errors when the code is not used for production purpose.
Instead, all torch.jit.script required classes and functions are tested in test_diar_utils.py.
Take a look at these tests for checking jit_script = [True/False] and cuda = [True/False] (testing total 4 combinations)
Also refactored some of the functions in online diarization
- replaced scipy LSA solver to NeMo LSA solver in online_clustering.py.
Added a couple of functions in der.py for online diarization DER calculation.
- replaced scipy LSA solver to NeMo LSA solver in der.py.

Collection: [ASR]

Changelog

nemo/collections/asr/metrics/der.py
: replaced scipy LSA solver to NeMo LSA solver in calculate_session_cpWER function.
: Added two functions for online diarization evaluations: get_partial_ref_labels and get_online_DER_stats.
nemo/collections/asr/models/online_diarizer.py
: Made _perform_online_clustering function simpler by moving get_reduced_mat and match_labels into online clustering function.
nemo/collections/asr/parts/utils/offline_clustering.py
: Added laplacian = laplacian.float().to(torch.device('cpu')) to avoid jit-scripted module uses GPU even when CPU is specified or vice-versa. This behavior is always tested/checked in test_diar_utils.py.
nemo/collections/asr/parts/utils/online_clustering.py
: replaced scipy LSA solver to NeMo LSA solver in get_lsa_speaker_mapping function.
: Modified the docstrings of update_speaker_history_buffer to make the example easier.
nemo/collections/asr/parts/utils/optimization_utils.py
: Fully torch-jit-scriptable, linear sum assignment problem solver class and function were added.
nemo/collections/asr/parts/utils/speaker_utils.py
: Removed @torch.jit.script decorators since this creates unnecessary warning messages and type related errors when used without scripting.
tests/collections/asr/test_diar_metrics.py
: Added unit-tests for the newly added function get_partial_ref_labels and get_online_DER_stats.
tests/collections/asr/test_diar_utils.py
: Added tests for offline clustering and online clustering for many different cases including:
[jit-script=True, cuda=True],
[jit-script=True, cuda=False],
[jit-script=False, cuda=True],
[jit-script=False, cuda=False] cases
which is using the torch-jit-scripted NeMo linear_sum_assignment function.

Usage

from nemo.collections.asr.parts.utils.optimization_utils import linear_sum_assignment
#An example cost matrix to be solved
cost_matrix = \
torch.tensor([[7, 6, 2], [6, 2, 1], [5, 6, 8]])
row_ind_nm, col_ind_nm = linear_sum_assignment(cost_matrix)

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.

…diarization Signed-off-by: Taejin Park <tango4j@gmail.com>

Signed-off-by: Taejin Park <tango4j@gmail.com>

for more information, see https://pre-commit.ci

Signed-off-by: Taejin Park <tango4j@gmail.com>

for more information, see https://pre-commit.ci

Signed-off-by: Taejin Park <tango4j@gmail.com>

…fix/clus_spk_util_jit

Signed-off-by: Taejin Park <tango4j@gmail.com>

for more information, see https://pre-commit.ci

Signed-off-by: Taejin Park <tango4j@gmail.com>

…fix/clus_spk_util_jit

nemo/collections/asr/models/online_diarizer.py

nemo/collections/asr/metrics/der.py

Signed-off-by: Taejin Park <tango4j@gmail.com>

for more information, see https://pre-commit.ci

Signed-off-by: Taejin Park <tango4j@gmail.com>

…fix/clus_spk_util_jit

nemo/collections/asr/metrics/der.py

Signed-off-by: Taejin Park <tango4j@gmail.com>

nemo/collections/asr/metrics/der.py

Signed-off-by: Taejin Park <tango4j@gmail.com>

for more information, see https://pre-commit.ci

Signed-off-by: Taejin Park <tango4j@gmail.com>

…fix/clus_spk_util_jit

for more information, see https://pre-commit.ci

nithinraok

minor review. Will do thorough review tomorrow.
Very neat improvement, need to understand better from my end.

nithinraok · 2023-04-07T04:27:48Z

nemo/collections/asr/parts/utils/offline_clustering.py

@@ -552,12 +550,13 @@ def eigDecompose(
            device = torch.cuda.current_device()
        laplacian = laplacian.float().to(device)
    else:
+        laplacian = laplacian.float().to(torch.device('cpu'))
        laplacian = laplacian.float()


why same operation twice?

nithinraok · 2023-04-07T04:28:45Z

nemo/collections/asr/parts/utils/offline_clustering.py

        laplacian = laplacian.float()
    lambdas, diffusion_map = eigh(laplacian)
    return lambdas, diffusion_map


-def eigValueSh(laplacian: torch.Tensor, cuda: bool, device: torch.device = torch.device('cpu')) -> torch.Tensor:
+def eigValueSh(laplacian: torch.Tensor, cuda: bool, device: torch.device) -> torch.Tensor:


why cuda and device? Isn't only one sufficient

This was added long back because there are users setting cuda=True but device=cpu.
This is adding some flexibility to avoid errors on such cases.
If we need to remove this, lt requires a speparate PR since this involves whole diarization pipeline.

nithinraok · 2023-04-07T04:29:18Z

nemo/collections/asr/parts/utils/offline_clustering.py

+        laplacian = laplacian.float().to(torch.device('cpu'))
        laplacian = laplacian.float()


same here. laplacian.float() twice

nithinraok · 2023-04-07T04:30:31Z

nemo/collections/asr/parts/utils/online_clustering.py

-    stacked = np.hstack((enc_P, enc_Q))
-    cost = -1 * linear_kernel(stacked.T)[spk_count:, :spk_count]
-    row_ind, col_ind = linear_sum_assignment(cost)
+    PandQ_list: List[int] = [int(x.item()) for x in PandQ]


minor: mentioning dtype in variable name need to be avoided

Makes sense since types are strictly annotated for jit script functions.
Fixed.

nithinraok · 2023-04-07T21:22:29Z

nemo/collections/asr/parts/utils/optimization_utils.py

+        marked (Tensor): 2D matrix containing the marked zeros.
+    """
+
+    def __init__(self, cost_matrix):


minor, mention the dtype of cost_matrix here. Isn;t it necessary for jit scripting?

If there is no type annotation, jit compiler think of it as torch.Tensor.
So in general if it is not torch.Tensor, type annotation is needed.
Added type annotations

nithinraok · 2023-04-07T21:29:24Z

nemo/collections/asr/parts/utils/optimization_utils.py

+    if cost_matrix.shape[1] < cost_matrix.shape[0]:
+        cost_matrix = cost_matrix.T
+        transposed = True
+    else:
+        transposed = False


why extra transposed variable, Use the same col < row condition below?

This followed the original implementation in scipy.
If we don't use transposed variable, we need to create another variable to indicate that foo = cost_matrix.shape[1] < cost_matrix.shape[0].

nithinraok · 2023-04-07T21:30:53Z

nemo/collections/asr/parts/utils/optimization_utils.py

+# Copyright (c) 2008 Brian M. Clapper <bmc@clapper.org>, Gael Varoquaux
+# Author: Brian M. Clapper, Gael Varoquaux
+# License: 3-clause BSD
+


Do we have only one optimization algorithm yet? Thinking if we should move other funcs to this file as well

I think we can add other algorithms below this. (I mentioned "Linear Sum Assignment solver")
The copyright in the beginning of the code is the convention in the most of the project so I followed

nithinraok · 2023-04-07T22:02:47Z

nemo/collections/asr/metrics/der.py

+    for label in ref_labels:
+        start, end, speaker = label.split()
+        start, end = float(start), float(end)
+        # If the current [start, end] interval is latching the last prediction time


latching -> matching

Changed the expression (Checked by Elena)

Signed-off-by: Taejin Park <tango4j@gmail.com>

for more information, see https://pre-commit.ci

Signed-off-by: Taejin Park <tango4j@gmail.com>

…fix/clus_spk_util_jit

nemo/collections/asr/parts/utils/online_clustering.py

@@ -31,67 +31,67 @@
 # https://arxiv.org/pdf/2003.02405.pdf and the implementation from
 # https://github.com/tango4j/Auto-Tuning-Spectral-Clustering.

-from typing import List, Tuple
+from typing import List, Set, Tuple


nithinraok

LGTM

…IA#6349) * [ASR] Add optimization utils for cpWER, diarization training, online diarization Signed-off-by: Taejin Park <tango4j@gmail.com> * Fixed GPU/CPU issues for clustering Signed-off-by: Taejin Park <tango4j@gmail.com> * Fixed unreachable state Signed-off-by: Taejin Park <tango4j@gmail.com> * resolved jit script compile error for lsa algorithm Signed-off-by: Taejin Park <tango4j@gmail.com> * Fixed errors and bugs, checked tests Signed-off-by: Taejin Park <tango4j@gmail.com> * Fixed docstrings Signed-off-by: Taejin Park <tango4j@gmail.com> * Update changes on test files Signed-off-by: Taejin Park <tango4j@gmail.com> * Refactored functions Signed-off-by: Taejin Park <tango4j@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Adding docstrings for the functions in der.py Signed-off-by: Taejin Park <tango4j@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed wrong docstrings in der.py Signed-off-by: Taejin Park <tango4j@gmail.com> * Fixed a wrong docstring Signed-off-by: Taejin Park <tango4j@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Changed np.array input to Tensor for LSA solver in der.py Signed-off-by: Taejin Park <tango4j@gmail.com> * Added code-QL issues and unit-tests for der.py functions Signed-off-by: Taejin Park <tango4j@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Removed print line in der.py Signed-off-by: Taejin Park <tango4j@gmail.com> * Fixed code QL redundant comparison Signed-off-by: Taejin Park <tango4j@gmail.com> * Fixed code QL issue Signed-off-by: Taejin Park <tango4j@gmail.com> * Added License for the reference code Signed-off-by: Taejin Park <tango4j@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added full license text of the original code Signed-off-by: Taejin Park <tango4j@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reflected comments Signed-off-by: Taejin Park <tango4j@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reflected review comments Signed-off-by: Taejin Park <tango4j@gmail.com> --------- Signed-off-by: Taejin Park <tango4j@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: hsiehjackson <c2hsieh@ucsd.edu>

tango4j added 9 commits March 24, 2023 09:51

[ASR] Add optimization utils for cpWER, diarization training, online …

d8c40f5

…diarization Signed-off-by: Taejin Park <tango4j@gmail.com>

Fixed GPU/CPU issues for clustering

c8884bf

Signed-off-by: Taejin Park <tango4j@gmail.com>

Fixed unreachable state

c083836

Signed-off-by: Taejin Park <tango4j@gmail.com>

resolved jit script compile error for lsa algorithm

8c801d2

Signed-off-by: Taejin Park <tango4j@gmail.com>

Fixed errors and bugs, checked tests

f238371

Signed-off-by: Taejin Park <tango4j@gmail.com>

Fixed docstrings

6464526

Signed-off-by: Taejin Park <tango4j@gmail.com>

Update changes on test files

bdf183d

Signed-off-by: Taejin Park <tango4j@gmail.com>

Merge remote-tracking branch 'origin' into fix/clus_spk_util_jit

53234d3

Refactored functions

4abc533

Signed-off-by: Taejin Park <tango4j@gmail.com>

github-actions bot added the ASR label Apr 3, 2023

[pre-commit.ci] auto fixes from pre-commit.com hooks

49d4e80

for more information, see https://pre-commit.ci

tango4j requested a review from nithinraok April 3, 2023 02:31

tango4j and others added 9 commits April 2, 2023 19:38

Adding docstrings for the functions in der.py

225f2f3

Signed-off-by: Taejin Park <tango4j@gmail.com>

Resolved conflict and adding docstrings to der.py

be65ec7

Signed-off-by: Taejin Park <tango4j@gmail.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

8ebeee3

for more information, see https://pre-commit.ci

Fixed wrong docstrings in der.py

154470b

Signed-off-by: Taejin Park <tango4j@gmail.com>

Merge branch 'fix/clus_spk_util_jit' of github.com:tango4j/NeMo into …

d9fe7b6

…fix/clus_spk_util_jit

Fixed a wrong docstring

891c1df

Signed-off-by: Taejin Park <tango4j@gmail.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

fe7102b

for more information, see https://pre-commit.ci

Changed np.array input to Tensor for LSA solver in der.py

a5b7807

Signed-off-by: Taejin Park <tango4j@gmail.com>

Merge branch 'fix/clus_spk_util_jit' of github.com:tango4j/NeMo into …

03be164

…fix/clus_spk_util_jit

github-advanced-security bot found potential problems Apr 3, 2023

View reviewed changes

nemo/collections/asr/models/online_diarizer.py Fixed Show resolved Hide resolved

nemo/collections/asr/metrics/der.py Fixed Show resolved Hide resolved

tango4j and others added 3 commits April 3, 2023 16:16

Added code-QL issues and unit-tests for der.py functions

d0dc2d3

Signed-off-by: Taejin Park <tango4j@gmail.com>

Merge branch 'main' into fix/clus_spk_util_jit

520ad6b

[pre-commit.ci] auto fixes from pre-commit.com hooks

60e6a06

for more information, see https://pre-commit.ci

tango4j marked this pull request as ready for review April 3, 2023 23:30

tango4j requested a review from fayejf April 3, 2023 23:36

tango4j marked this pull request as draft April 3, 2023 23:37

tango4j added 2 commits April 3, 2023 16:39

Removed print line in der.py

d996a5d

Signed-off-by: Taejin Park <tango4j@gmail.com>

Merge branch 'fix/clus_spk_util_jit' of github.com:tango4j/NeMo into …

1e812f8

…fix/clus_spk_util_jit

github-advanced-security bot found potential problems Apr 4, 2023

View reviewed changes

nemo/collections/asr/metrics/der.py Fixed Show fixed Hide fixed

Fixed code QL redundant comparison

b342716

Signed-off-by: Taejin Park <tango4j@gmail.com>

github-advanced-security bot found potential problems Apr 4, 2023

View reviewed changes

nemo/collections/asr/metrics/der.py Fixed Show fixed Hide fixed

tango4j marked this pull request as ready for review April 4, 2023 18:05

tango4j and others added 6 commits April 4, 2023 14:10

Fixed code QL issue

712b4bd

Signed-off-by: Taejin Park <tango4j@gmail.com>

Added License for the reference code

0057e4d

Signed-off-by: Taejin Park <tango4j@gmail.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

c13a04f

for more information, see https://pre-commit.ci

Added full license text of the original code

4647504

Signed-off-by: Taejin Park <tango4j@gmail.com>

Merge branch 'fix/clus_spk_util_jit' of github.com:tango4j/NeMo into …

db7d3fa

…fix/clus_spk_util_jit

[pre-commit.ci] auto fixes from pre-commit.com hooks

6f08917

for more information, see https://pre-commit.ci

nithinraok reviewed Apr 7, 2023

View reviewed changes

tango4j and others added 6 commits April 12, 2023 17:42

Merge remote-tracking branch 'origin' into fix/clus_spk_util_jit

5184600

Reflected comments

cd0394d

Signed-off-by: Taejin Park <tango4j@gmail.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

5a634df

for more information, see https://pre-commit.ci

Merge branch 'main' into fix/clus_spk_util_jit

09ecfa4

Reflected review comments

443b6af

Signed-off-by: Taejin Park <tango4j@gmail.com>

Merge branch 'fix/clus_spk_util_jit' of github.com:tango4j/NeMo into …

73861cf

…fix/clus_spk_util_jit

github-advanced-security bot found potential problems Apr 14, 2023

View reviewed changes

nithinraok approved these changes Apr 14, 2023

View reviewed changes

tango4j merged commit ae55b52 into NVIDIA:main Apr 14, 2023

tango4j deleted the fix/clus_spk_util_jit branch December 6, 2023 21:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ASR] Add optimization util for linear sum assignment algorithm #6349

[ASR] Add optimization util for linear sum assignment algorithm #6349

tango4j commented Apr 3, 2023 •

edited

Loading

nithinraok left a comment

nithinraok Apr 7, 2023

tango4j Apr 13, 2023

nithinraok Apr 7, 2023

tango4j Apr 13, 2023

nithinraok Apr 7, 2023

tango4j Apr 13, 2023

nithinraok Apr 7, 2023

tango4j Apr 14, 2023

nithinraok Apr 7, 2023

tango4j Apr 13, 2023 •

edited

Loading

nithinraok Apr 14, 2023

nithinraok Apr 7, 2023

tango4j Apr 13, 2023

nithinraok Apr 7, 2023

tango4j Apr 13, 2023 •

edited

Loading

nithinraok Apr 7, 2023

tango4j Apr 13, 2023

nithinraok left a comment

		laplacian = laplacian.float().to(torch.device('cpu'))
		laplacian = laplacian.float()

[ASR] Add optimization util for linear sum assignment algorithm #6349

[ASR] Add optimization util for linear sum assignment algorithm #6349

Conversation

tango4j commented Apr 3, 2023 • edited Loading

What does this PR do ?

Changelog

Usage

Before your PR is "Ready for review"

Who can review?

nithinraok left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tango4j Apr 13, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tango4j Apr 13, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nithinraok left a comment

Choose a reason for hiding this comment

tango4j commented Apr 3, 2023 •

edited

Loading

tango4j Apr 13, 2023 •

edited

Loading

tango4j Apr 13, 2023 •

edited

Loading