Search the installed NeMo package and the local project for all available auxiliary losses, multi-objective training options, and loss functions that can be used with EncDecHybridRNNTCTCBPEModel or similar ASR models.

Look in:
1. The installed nemo package: search for loss classes, auxiliary objectives, multi-task losses in /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/
2. Specifically look at:
   - nemo/collections/asr/losses/ directory
   - nemo/collections/asr/models/ for any multi-objective or hybrid models
   - Any RNNT/TDT/CTC/transducer loss variants
   - Any language ID losses, contrastive losses, or self-supervised objectives
3. Also check if there are any auxiliary loss hooks or mixins

Report all loss classes found, their purpose, and whether they could be added as auxiliary objectives. Be thorough - read the actual loss files.


I'll conduct a thorough search of the NeMo package and the local project to identify all available auxiliary losses, multi-objective training options, and loss functions.


Bash

find /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr -type f -name "*.py" | head -20


Let me check the losses directory specifically:


total 80
drwxrwxr-x 4 ubuntu ubuntu  4096 Feb 21 08:40 .
drwxrwxr-x 9 ubuntu ubuntu  4096 Feb 21 08:40 ..
-rw-rw-r-- 1 ubuntu ubuntu  2362 Feb 21 08:40 angularloss.py
-rw-rw-r-- 1 ubuntu ubuntu  5844 Feb 21 08:40 bce_loss.py
-rw-rw-r-- 1 ubuntu ubuntu  3301 Feb 21 08:40 ctc.py
-rw-rw-r-- 1 ubuntu ubuntu  1154 Feb 21 08:40 __init__.py
-rw-rw-r-- 1 ubuntu ubuntu  7727 Feb 21 08:40 lattice_losses.py
drwxrwxr-x 2 ubuntu ubuntu  4096 Feb 21 08:40 __pycache__
-rw-rw-r-- 1 ubuntu ubuntu 20243 Feb 21 08:40 rnnt.py
-rw-rw-r-- 1 ubuntu ubuntu 16065 Feb 21 08:40 rnnt_pytorch.py
drwxrwxr-x 3 ubuntu ubuntu  4096 Feb 21 08:40 ssl_losses


Great! Now let me search through the models directory and get all the loss files:


1	# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	from nemo.collections.asr.losses.angularloss import AngularSoftmaxLoss
16	from nemo.collections.asr.losses.bce_loss import BCELoss
17	from nemo.collections.asr.losses.ctc import CTCLoss
18	from nemo.collections.asr.losses.lattice_losses import LatticeLoss
19	from nemo.collections.asr.losses.ssl_losses.contrastive import ContrastiveLoss
20	from nemo.collections.asr.losses.ssl_losses.ctc import CTCLossForSSL
21	from nemo.collections.asr.losses.ssl_losses.mlm import MLMLoss, MultiMLMLoss
22	from nemo.collections.asr.losses.ssl_losses.rnnt import RNNTLossForSSL
23


total 48
drwxrwxr-x 3 ubuntu ubuntu  4096 Feb 21 08:40 .
drwxrwxr-x 4 ubuntu ubuntu  4096 Feb 21 08:40 ..
-rw-rw-r-- 1 ubuntu ubuntu 13416 Feb 21 08:40 contrastive.py
-rw-rw-r-- 1 ubuntu ubuntu  2155 Feb 21 08:40 ctc.py
-rw-rw-r-- 1 ubuntu ubuntu   690 Feb 21 08:40 __init__.py
-rw-rw-r-- 1 ubuntu ubuntu  5015 Feb 21 08:40 mlm.py
drwxrwxr-x 2 ubuntu ubuntu  4096 Feb 21 08:40 __pycache__
-rw-rw-r-- 1 ubuntu ubuntu  2084 Feb 21 08:40 rnnt.py


find /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models -type f -name "*.py" | grep -E "(hybrid|rnnt|ctc|transducer)" | head -20


Now let me read all the loss files:


1	# ! /usr/bin/python
2	# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	#     http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	import torch
17	from torch import nn
18
19	from nemo.core.classes import Serialization, Typing, typecheck
20	from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
21
22	__all__ = ['CTCLoss']
23
24
25	class CTCLoss(nn.CTCLoss, Serialization, Typing):
26	    @property
27	    def input_types(self):
28	        """Input types definitions for CTCLoss.
29	        """
30	        return {
31	            "log_probs": NeuralType(('B', 'T', 'D'), LogprobsType()),
32	            "targets": NeuralType(('B', 'T'), LabelsType()),
33	            "input_lengths": NeuralType(tuple('B'), LengthsType()),
34	            "target_lengths": NeuralType(tuple('B'), LengthsType()),
35	        }
36
37	    @property
38	    def output_types(self):
39	        """Output types definitions for CTCLoss.
40	        loss:
41	            NeuralType(None)
42	        """
43	        return {"loss": NeuralType(elements_type=LossType())}
44
45	    def __init__(self, num_classes, zero_infinity=False, reduction='mean_batch'):
46	        self._blank = num_classes
47	        # Don't forget to properly call base constructor
48	        if reduction not in ['none', 'mean', 'sum', 'mean_batch', 'mean_volume']:
49	            raise ValueError('`reduction` must be one of [mean, sum, mean_batch, mean_volume]')
50
51	        self.config_reduction = reduction
52	        if reduction == 'mean_batch' or reduction == 'mean_volume':
53	            ctc_reduction = 'none'
54	            self._apply_reduction = True
55	        elif reduction in ['sum', 'mean', 'none']:
56	            ctc_reduction = reduction
57	            self._apply_reduction = False
58	        super().__init__(blank=self._blank, reduction=ctc_reduction, zero_infinity=zero_infinity)
59
60	    def reduce(self, losses, target_lengths):
61	        if self.config_reduction == 'mean_batch':
62	            losses = losses.mean()  # global batch size average
63	        elif self.config_reduction == 'mean_volume':
64	            losses = losses.sum() / target_lengths.sum()  # same as above but longer samples weigh more
65
66	        return losses
67
68	    @typecheck()
69	    def forward(self, log_probs, targets, input_lengths, target_lengths):
70	        # override forward implementation
71	        # custom logic, if necessary
72	        input_lengths = input_lengths.long()
73	        target_lengths = target_lengths.long()
74	        targets = targets.long()
75	        # here we transpose because we expect [B, T, D] while PyTorch assumes [T, B, D]
76	        log_probs = log_probs.transpose(1, 0)
77	        loss = super().forward(
78	            log_probs=log_probs, targets=targets, input_lengths=input_lengths, target_lengths=target_lengths
79	        )
80	        if self._apply_reduction:
81	            loss = self.reduce(loss, target_lengths)
82	        return loss
83


1	# ! /usr/bin/python
2	# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	#     http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	# Copyright 2018-2019, Mingkun Huang
17	#
18	# Licensed under the Apache License, Version 2.0 (the "License");
19	# you may not use this file except in compliance with the License.
20	# You may obtain a copy of the License at
21	#
22	#    http://www.apache.org/licenses/LICENSE-2.0
23	#
24	# Unless required by applicable law or agreed to in writing, software
25	# distributed under the License is distributed on an "AS IS" BASIS,
26	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
27	# See the License for the specific language governing permissions and
28	# limitations under the License.
29
30	import inspect
31	import operator
32	from dataclasses import dataclass
33	from typing import Any, Callable, Dict, List, Optional, Set
34
35	import torch
36	from omegaconf import DictConfig, OmegaConf
37
38	from nemo.collections.asr.losses.rnnt_pytorch import MultiblankRNNTLossPytorch, RNNTLossPytorch, TDTLossPytorch
39	from nemo.core.classes import Loss, typecheck
40	from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
41	from nemo.core.utils import numba_utils
42	from nemo.core.utils.k2_utils import K2_INSTALLATION_MESSAGE
43	from nemo.core.utils.numba_utils import NUMBA_INSTALLATION_MESSAGE
44	from nemo.utils import logging, logging_mode, model_utils
45
46	try:
47	    import warprnnt_pytorch as warprnnt
48
49	    WARP_RNNT_AVAILABLE = True
50	except (ImportError, ModuleNotFoundError):
51	    WARP_RNNT_AVAILABLE = False
52
53	try:
54	    from nemo.collections.asr.parts.numba.rnnt_loss import MultiblankRNNTLossNumba, RNNTLossNumba, TDTLossNumba
55
56	    NUMBA_RNNT_AVAILABLE = True
57	except (ImportError, ModuleNotFoundError):
58	    NUMBA_RNNT_AVAILABLE = False
59
60	try:
61	    from nemo.collections.asr.parts.k2.graph_transducer import GraphRnntLoss
62	    from nemo.collections.asr.parts.k2.w_transducer import GraphWTransducerLoss
63
64	    K2_AVAILABLE = True
65	except (ImportError, ModuleNotFoundError):
66	    K2_AVAILABLE = False
67
68	WARP_RNNT_INSTALLATION_MESSAGE = (
69	    "Could not import `warprnnt_pytorch`.\n"
70	    "Please visit https://github.com/HawkAaron/warp-transducer "
71	    "and follow the steps in the readme to build and install the "
72	    "pytorch bindings for RNNT Loss, or use the provided docker "
73	    "container that supports RNN-T loss."
74	)
75
76
77	@dataclass
78	class RNNTLossConfig:
79	    loss_name: str
80	    lib_name: str
81	    is_available: bool = False
82	    installation_msg: str = ""
83	    min_version: Optional[str] = None
84	    force_float32: bool = True  # default True for now for all losses except graph-based
85
86
87	# Resolved list of available RNNT losses
88	RNNT_LOSS_RESOLVER = {
89	    "warprnnt": RNNTLossConfig(
90	        loss_name="warprnnt",
91	        lib_name="warprnnt_pytorch",
92	        is_available=WARP_RNNT_AVAILABLE,
93	        installation_msg=WARP_RNNT_INSTALLATION_MESSAGE,
94	        force_float32=True,
95	    ),
96	    "warprnnt_numba": RNNTLossConfig(
97	        loss_name="warprnnt_numba",
98	        lib_name="numba",
99	        min_version='0.53.0',
100	        is_available=NUMBA_RNNT_AVAILABLE,
101	        installation_msg=NUMBA_INSTALLATION_MESSAGE,
102	        force_float32=False,  # This is only temporarily false, will be dynamically updated during resolution
103	    ),
104	    "pytorch": RNNTLossConfig(
105	        loss_name="pytorch",
106	        lib_name="torch",
107	        min_version='0.0',
108	        is_available=True,
109	        installation_msg="Pure Pytorch implementation of RNN-T loss. Slow and for debugging purposes only.",
110	        force_float32=True,
111	    ),
112	    "multiblank_rnnt": RNNTLossConfig(
113	        loss_name="multiblank_rnnt",
114	        lib_name="numba",
115	        min_version='0.53.0',
116	        is_available=NUMBA_RNNT_AVAILABLE,
117	        installation_msg=NUMBA_INSTALLATION_MESSAGE,
118	        force_float32=True,
119	    ),
120	    "multiblank_rnnt_pytorch": RNNTLossConfig(
121	        loss_name="pytorch",
122	        lib_name="torch",
123	        min_version='0.0',
124	        is_available=True,
125	        installation_msg="Pure Pytorch implementation of Multiblank RNN-T loss. Slow and for debugging purposes only.",
126	        force_float32=True,
127	    ),
128	    "graph_w_transducer": RNNTLossConfig(
129	        loss_name="graph_w_transducer",
130	        lib_name="k2",
131	        is_available=K2_AVAILABLE,
132	        installation_msg=K2_INSTALLATION_MESSAGE,
133	        force_float32=False,
134	    ),
135	    "graph_rnnt": RNNTLossConfig(
136	        loss_name="graph_rnnt",
137	        lib_name="k2",
138	        is_available=K2_AVAILABLE,
139	        installation_msg=K2_INSTALLATION_MESSAGE,
140	        force_float32=False,
141	    ),
142	    "tdt": RNNTLossConfig(
143	        loss_name="tdt",
144	        lib_name="numba",
145	        min_version='0.53.0',
146	        is_available=NUMBA_RNNT_AVAILABLE,
147	        installation_msg=NUMBA_INSTALLATION_MESSAGE,
148	    ),
149	    "tdt_pytorch": RNNTLossConfig(
150	        loss_name="tdt_pytorch",
151	        lib_name="torch",
152	        min_version='0.0',
153	        is_available=True,
154	        installation_msg="Pure Pytorch implementation of TDT loss. Slow and for debugging purposes only.",
155	    ),
156	}
157
158	RNNT_LOSS_RESOLVER['default'] = RNNT_LOSS_RESOLVER['warprnnt_numba']
159
160
161	def _warn_unused_additional_kwargs(loss_name, kwargs):
162	    if len(kwargs) > 0:
163	        logging.warning(
164	            f"Loss function `{loss_name}` was provided with following additional kwargs,\n"
165	            f"however they were ignored as it is unused.\n"
166	            f"{kwargs}"
167	        )
168
169
170	def _clean_kwargs(
171	    loss_name: str, kwargs: Optional[Dict[str, Any]], init_method: Callable, ignore_params: Optional[Set[str]] = None
172	) -> Dict[str, Any]:
173	    """
174	    Cleans kwargs for the given loss function. Warn if there are unused kwargs.
175
176	    Args:
177	        loss_name: name of the loss function
178	        kwargs: kwargs to clean
179	        init_method: LossClass.__init__ method
180	        ignore_params: set of argument names for init_method to ignore
181
182	    Returns:
183	        only used kwargs for the given `init_method`
184	    """
185	    if not kwargs:
186	        return {}
187	    init_params = set(inspect.signature(init_method).parameters.keys()) - {"self"}
188	    if ignore_params is not None:
189	        init_params -= ignore_params
190	    unused_kwargs = dict()
191	    used_kwargs = dict()
192	    for key, value in kwargs.items():
193	        if key not in init_params:
194	            unused_kwargs[key] = value
195	        else:
196	            used_kwargs[key] = value
197	    if len(unused_kwargs) > 0:
198	        _warn_unused_additional_kwargs(loss_name, unused_kwargs)
199	    return used_kwargs
200
201
202	def resolve_rnnt_default_loss_name() -> str:
203	    return RNNT_LOSS_RESOLVER['default'].loss_name
204
205
206	def resolve_rnnt_loss(loss_name: str, blank_idx: int, loss_kwargs: dict = None) -> torch.nn.Module:
207	    loss_function_names = list(RNNT_LOSS_RESOLVER.keys())
208
209	    if loss_name not in loss_function_names:
210	        raise ValueError(
211	            f"Provided `loss_name` {loss_name} not in list of available RNNT losses \n" f"{loss_function_names}"
212	        )
213
214	    all_available_losses = {name: config for name, config in RNNT_LOSS_RESOLVER.items() if config.is_available}
215
216	    loss_config = RNNT_LOSS_RESOLVER[loss_name]  # type: RNNTLossConfig
217
218	    # Re-raise import error with installation message
219	    if not loss_config.is_available:
220	        msg = (
221	            f"Installed RNNT losses are : {list(all_available_losses.keys())}.\n"
222	            f"****************************************************************\n"
223	            f"To install the selected loss function, please follow the steps below:\n"
224	            f"{loss_config.installation_msg}"
225	        )
226	        raise ImportError(msg)
227
228	    # Library version check
229	    if loss_config.min_version is not None:
230	        ver_matched, msg = model_utils.check_lib_version(
231	            loss_config.lib_name, checked_version=loss_config.min_version, operator=operator.ge
232	        )
233
234	        if ver_matched is False:
235	            msg = (
236	                f"{msg}\n"
237	                f"****************************************************************\n"
238	                f"To update the selected loss function, please follow the steps below:\n"
239	                f"{loss_config.installation_msg}"
240	            )
241	            raise RuntimeError(msg)
242
243	    # Resolve loss functions sequentially
244	    loss_kwargs = {} if loss_kwargs is None else loss_kwargs
245
246	    if isinstance(loss_kwargs, DictConfig):
247	        loss_kwargs = OmegaConf.to_container(loss_kwargs, resolve=True)
248
249	    # Get actual loss name for `default`
250	    if loss_name == 'default':
251	        loss_name = loss_config.loss_name
252
253	    """
254	    Resolve RNNT loss functions
255	    """
256	    if loss_name == 'warprnnt':
257	        loss_func = warprnnt.RNNTLoss(blank=blank_idx, reduction='none')
258	        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
259
260	    elif loss_name == 'warprnnt_numba':
261	        # Update loss config's forced float32 flag if set to None
262	        loss_config.force_float32 = not numba_utils.is_numba_cuda_fp16_supported()
263
264	        fastemit_lambda = loss_kwargs.pop('fastemit_lambda', 0.0)
265	        clamp = loss_kwargs.pop('clamp', -1.0)
266	        loss_func = RNNTLossNumba(blank=blank_idx, reduction='none', fastemit_lambda=fastemit_lambda, clamp=clamp)
267	        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
268
269	    elif loss_name == 'pytorch':
270	        loss_func = RNNTLossPytorch(blank=blank_idx, reduction='none')
271	        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
272
273	    elif loss_name == 'multiblank_rnnt':
274	        fastemit_lambda = loss_kwargs.pop('fastemit_lambda', 0.0)
275	        clamp = loss_kwargs.pop('clamp', -1.0)
276	        big_blank_durations = loss_kwargs.pop('big_blank_durations', None)
277	        sigma = loss_kwargs.pop('sigma', 0.0)
278	        loss_func = MultiblankRNNTLossNumba(
279	            blank=blank_idx,
280	            big_blank_durations=big_blank_durations,
281	            reduction='none',
282	            fastemit_lambda=fastemit_lambda,
283	            clamp=clamp,
284	            sigma=sigma,
285	        )
286	        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
287
288	    elif loss_name == 'multiblank_rnnt_pytorch':
289	        big_blank_durations = loss_kwargs.pop('big_blank_durations', None)
290	        sigma = loss_kwargs.pop('sigma', 0.0)
291	        loss_func = MultiblankRNNTLossPytorch(
292	            blank=blank_idx, big_blank_durations=big_blank_durations, reduction='none', sigma=sigma
293	        )
294	        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
295
296	    elif loss_name == 'tdt':
297	        fastemit_lambda = loss_kwargs.pop('fastemit_lambda', 0.0)
298	        clamp = loss_kwargs.pop('clamp', -1.0)
299	        durations = loss_kwargs.pop('durations', None)
300	        sigma = loss_kwargs.pop('sigma', 0.0)
301	        omega = loss_kwargs.pop('omega', 0.0)
302	        loss_func = TDTLossNumba(
303	            blank=blank_idx,
304	            durations=durations,
305	            reduction='none',
306	            fastemit_lambda=fastemit_lambda,
307	            clamp=clamp,
308	            sigma=sigma,
309	            omega=omega,
310	        )
311	        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
312
313	    elif loss_name == 'tdt_pytorch':
314	        durations = loss_kwargs.pop('durations', None)
315	        sigma = loss_kwargs.pop('sigma', 0.0)
316	        loss_func = TDTLossPytorch(blank=blank_idx, durations=durations, reduction='none', sigma=sigma)
317	        _warn_unused_additional_kwargs(loss_name, loss_kwargs)
318
319	    elif loss_name == "graph_rnnt":
320	        loss_kwargs = _clean_kwargs(loss_name, loss_kwargs, GraphRnntLoss.__init__, ignore_params={"blank"})
321	        loss_func = GraphRnntLoss(blank=blank_idx, **loss_kwargs)
322	    elif loss_name == "graph_w_transducer":
323	        loss_kwargs = _clean_kwargs(loss_name, loss_kwargs, GraphWTransducerLoss.__init__, ignore_params={"blank"})
324	        loss_func = GraphWTransducerLoss(blank=blank_idx, **loss_kwargs)
325	    else:
326	        raise ValueError(
327	            f"Invalid value of `loss_name`: {loss_name}. Allowed loss names are :" f"{loss_function_names}"
328	        )
329
330	    return loss_func
331
332
333	class RNNTLoss(Loss):
334	    @property
335	    def input_types(self):
336	        """Input types definitions for CTCLoss.
337	        """
338	        return {
339	            "log_probs": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
340	            "targets": NeuralType(('B', 'T'), LabelsType()),
341	            "input_lengths": NeuralType(tuple('B'), LengthsType()),
342	            "target_lengths": NeuralType(tuple('B'), LengthsType()),
343	        }
344
345	    @property
346	    def output_types(self):
347	        """Output types definitions for CTCLoss.
348	        loss:
349	            NeuralType(None)
350	        """
351	        return {"loss": NeuralType(elements_type=LossType())}
352
353	    def __init__(self, num_classes, reduction: str = 'mean_batch', loss_name: str = "default", loss_kwargs=None):
354	        """
355	        RNN-T Loss function based on https://github.com/HawkAaron/warp-transducer.
356	        Optionally, can utilize a numba implementation of the same loss without having to compile the loss,
357	        albiet there is a small speed penalty for JIT numba compile.
358
359	        Note:
360	            Requires Numba 0.53.0 or later to be installed to use this loss function.
361
362	        Losses can be selected via the config, and optionally be passed keyword arguments as follows.
363
364	        Examples:
365	            .. code-block:: yaml
366
367	                model:  # RNNT Model config
368	                    ...
369	                    loss:
370	                        loss_name: "warprnnt_numba"
371	                        warprnnt_numba_kwargs:
372	                            fastemit_lambda: 0.0
373
374	        Warning:
375	            In the case that GPU memory is exhausted in order to compute RNNTLoss, it might cause
376	            a core dump at the cuda level with the following error message.
377
378	            ```
379	                ...
380	                costs = costs.to(acts.device)
381	            RuntimeError: CUDA error: an illegal memory access was encountered
382	            terminate called after throwing an instance of 'c10::Error'
383	            ```
384
385	            Please kill all remaining python processes after this point, and use a smaller batch size
386	            for train, validation and test sets so that CUDA memory is not exhausted.
387
388	        Args:
389	            num_classes: Number of target classes for the joint network to predict.
390	                In all cases (conventional RNNT, multi-blank RNNT, and TDT model), this equals the token-id
391	                for the standard "blank" symbol. In particular, say V is the number of non-blank tokens in
392	                the vocabulary, then in the case of,
393	                standard RNNT: num_classes = V
394	                multiblank RNNT: num_classes = V + number-big-blanks (since we store big-blanks before
395	                                 standard blank, and the standard blank is the last symbol in the vocab)
396	                TDT: num_classes = V. Note, V here does not include any of the "duration outputs".
397
398	            reduction: Type of reduction to perform on loss. Possible values are
399	                `mean_batch`, 'mean_volume`, `mean`, `sum` or None.
400	                `None` will return a torch vector comprising the individual loss values of the batch.
401	                `mean_batch` will average the losses in the batch
402	                `mean` will divide each loss by the target length and then average
403	                `mean_volume` will add up all the losses and divide by sum of target lengths
404
405	            loss_name: String that is resolved into an RNNT loss function. Available list of losses
406	                is ininitialized in `RNNT_LOSS_RESOLVER` dictionary.
407
408	            loss_kwargs: Optional Dict of (str, value) pairs that are passed to the instantiated loss
409	                function.
410	        """
411	        super(RNNTLoss, self).__init__()
412
413	        if reduction not in [None, 'mean', 'sum', 'mean_batch', 'mean_volume']:
414	            raise ValueError('`reduction` must be one of [mean, sum, mean_batch, mean_volume]')
415
416	        self._blank = num_classes
417	        self.reduction = reduction
418	        self._loss = resolve_rnnt_loss(loss_name, blank_idx=self._blank, loss_kwargs=loss_kwargs)
419	        self._force_float32 = RNNT_LOSS_RESOLVER[loss_name].force_float32
420	        self._fp16_compat_checked = False
421
422	    def reduce(self, losses, target_lengths):
423
424	        if isinstance(losses, List):
425	            losses = torch.cat(losses, 0)
426	            target_lengths = torch.cat(target_lengths, 0)
427
428	        if self.reduction == 'mean_batch':
429	            losses = losses.mean()  # global batch size average
430	        elif self.reduction == 'mean':
431	            losses = torch.div(losses, target_lengths).mean()
432	        elif self.reduction == 'sum':
433	            losses = losses.sum()
434	        elif self.reduction == 'mean_volume':
435	            losses = losses.sum() / target_lengths.sum()  # same as above but longer samples weigh more
436
437	        return losses
438
439	    @typecheck()
440	    def forward(self, log_probs, targets, input_lengths, target_lengths):
441	        # Cast to int 64
442	        targets = targets.long()
443	        input_lengths = input_lengths.long()
444	        target_lengths = target_lengths.long()
445
446	        max_logit_len = input_lengths.max()
447	        max_targets_len = target_lengths.max()
448
449	        # Force cast joint to float32
450	        if not self._force_float32 and numba_utils.is_numba_cuda_fp16_supported():
451	            # Execute the kernel in fp16
452	            pass
453	        elif self._force_float32 and log_probs.dtype != torch.float32:
454	            # Log just once if fp16 tensor was passed and fp16 Numba CUDA loss could not be used.
455	            if log_probs.dtype == torch.float16 and not self._fp16_compat_checked:
456	                _, reason = numba_utils.is_numba_cuda_fp16_supported(return_reason=True)
457	                logging.warning(
458	                    f"Provided RNNT Joint tensor is of dtype {log_probs.dtype}, but RNNT loss could not be calculated "
459	                    f"in fp16 due to following reason stated below. Loss will be calculated in fp32. \n\n"
460	                    f"{reason}",
461	                    mode=logging_mode.ONCE,
462	                )
463	                self._fp16_compat_checked = True
464
465	            # Upcast the activation tensor and compute loss and grads in fp32
466	            logits_orig = log_probs
467	            log_probs = log_probs.float()
468	            del logits_orig  # save memory *before* computing the loss
469
470	        # Ensure that shape mismatch does not occur due to padding
471	        # Due to padding and subsequent downsampling, it may be possible that
472	        # max sequence length computed does not match the actual max sequence length
473	        # of the log_probs tensor, therefore we increment the input_lengths by the difference.
474	        # This difference is generally small.
475	        if log_probs.shape[1] != max_logit_len:
476	            log_probs = log_probs.narrow(dim=1, start=0, length=max_logit_len).contiguous()
477
478	        # Reduce transcript length to correct alignment if additional padding was applied.
479	        # Transcript: [B, L] -> [B, L']; If L' < L
480	        if not targets.is_contiguous():
481	            targets = targets.contiguous()
482
483	        if targets.shape[1] != max_targets_len:
484	            targets = targets.narrow(dim=1, start=0, length=max_targets_len).contiguous()
485
486	        # Temporarily override loss reduction
487	        loss_reduction = self._loss.reduction
488	        self._loss.reduction = None
489
490	        # Compute RNNT loss
491	        loss = self._loss(acts=log_probs, labels=targets, act_lens=input_lengths, label_lens=target_lengths)
492
493	        # Loss reduction can be dynamic, so reset it after call
494	        self._loss.reduction = loss_reduction
495
496	        # reduce here using our own reduction function
497	        if self.reduction is not None:
498	            loss = self.reduce(loss, target_lengths)
499
500	        # del new variables that may have been created
501	        del (
502	            log_probs,
503	            targets,
504	            input_lengths,
505	            target_lengths,
506	        )
507
508	        return loss
509


1	# ! /usr/bin/python
2	# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	#     http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	from typing import List
17
18	import torch
19
20	from nemo.core.classes import Loss
21	from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
22
23
24	class RNNTLossPytorch(Loss):
25	    @property
26	    def input_types(self):
27	        """Input types definitions for CTCLoss.
28	        """
29	        return {
30	            "acts": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
31	            "labels": NeuralType(('B', 'T'), LabelsType()),
32	            "act_lens": NeuralType(tuple('B'), LengthsType()),
33	            "label_lens": NeuralType(tuple('B'), LengthsType()),
34	        }
35
36	    @property
37	    def output_types(self):
38	        """Output types definitions for CTCLoss.
39	        loss:
40	            NeuralType(None)
41	        """
42	        return {"loss": NeuralType(elements_type=LossType())}
43
44	    def __init__(self, blank, reduction):
45	        super().__init__()
46	        self.blank = blank
47	        self.reduction = reduction
48
49	    def forward(self, acts, labels, act_lens, label_lens):
50	        # CPU patch for FP16
51	        if not acts.is_cuda and acts.dtype == torch.float16:
52	            acts = acts.float()
53
54	        acts = torch.log_softmax(acts, -1)
55
56	        forward_logprob = self.compute_forward_prob(acts, labels, act_lens, label_lens)
57	        losses = -forward_logprob
58	        if self.reduction == 'mean_batch':
59	            losses = losses.mean()  # global batch size average
60	        elif self.reduction == 'mean':
61	            losses = torch.div(losses, label_lens).mean()
62	        elif self.reduction == 'sum':
63	            losses = losses.sum()
64	        elif self.reduction == 'mean_volume':
65	            losses = losses.sum() / label_lens.sum()  # same as above but longer samples weigh more
66
67	        return losses
68
69	    def compute_forward_prob(self, acts, labels, act_lens, label_lens):
70	        B, T, U, _ = acts.shape
71
72	        log_alpha = torch.zeros(B, T, U)
73	        log_alpha = log_alpha.to(acts.device)
74
75	        for t in range(T):
76	            for u in range(U):
77	                if u == 0:
78	                    if t == 0:
79	                        # this is the base case: (t=0, u=0) with log-alpha = 0.
80	                        log_alpha[:, t, u] = 0.0
81	                    else:
82	                        # this is case for (t = 0, u > 0), reached by (t, u - 1)
83	                        # emitting a blank symbol.
84	                        log_alpha[:, t, u] = log_alpha[:, t - 1, u] + acts[:, t - 1, 0, self.blank]
85	                else:
86	                    if t == 0:
87	                        # in case of (u > 0, t = 0), this is only reached from
88	                        # (t, u - 1) with a label emission.
89	                        gathered = torch.gather(
90	                            acts[:, t, u - 1], dim=1, index=labels[:, u - 1].view(-1, 1).type(torch.int64)
91	                        ).reshape(-1)
92	                        log_alpha[:, t, u] = log_alpha[:, t, u - 1] + gathered.to(log_alpha.device)
93	                    else:
94	                        # here both t and u are > 0, this state is reachable
95	                        # with two possibilities: (t - 1, u) with a blank emission
96	                        # or (t, u - 1) with a label emission.
97	                        log_alpha[:, t, u] = torch.logsumexp(
98	                            torch.stack(
99	                                [
100	                                    log_alpha[:, t - 1, u] + acts[:, t - 1, u, self.blank],
101	                                    log_alpha[:, t, u - 1]
102	                                    + torch.gather(
103	                                        acts[:, t, u - 1], dim=1, index=labels[:, u - 1].view(-1, 1).type(torch.int64)
104	                                    ).reshape(-1),
105	                                ]
106	                            ),
107	                            dim=0,
108	                        )
109
110	        log_probs = []
111	        for b in range(B):
112	            # here we need to add the final blank emission weights.
113	            to_append = (
114	                log_alpha[b, act_lens[b] - 1, label_lens[b]] + acts[b, act_lens[b] - 1, label_lens[b], self.blank]
115	            )
116	            log_probs.append(to_append)
117	        log_prob = torch.stack(log_probs)
118
119	        return log_prob
120
121
122	class TDTLossPytorch(Loss):
123	    """
124	    Pure Python implementation of TDT loss (https://arxiv.org/pdf/2304.06795.pdf)
125	    """
126
127	    @property
128	    def input_types(self):
129	        """Input types definitions for CTCLoss.
130	        """
131	        return {
132	            "acts": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
133	            "labels": NeuralType(('B', 'T'), LabelsType()),
134	            "act_lens": NeuralType(tuple('B'), LengthsType()),
135	            "label_lens": NeuralType(tuple('B'), LengthsType()),
136	        }
137
138	    @property
139	    def output_types(self):
140	        """Output types definitions for CTCLoss.
141	        loss:
142	            NeuralType(None)
143	        """
144	        return {"loss": NeuralType(elements_type=LossType())}
145
146	    def __init__(self, blank: int, durations: List[int] = [], reduction: str = 'sum', sigma: float = 0.0):
147	        super().__init__()
148	        self.blank = blank
149	        self.durations = durations
150	        self.n_durations = len(durations)
151	        self.reduction = reduction
152	        self.sigma = sigma
153
154	    def forward(self, acts, labels, act_lens, label_lens):
155	        label_acts = acts[:, :, :, : -self.n_durations]
156	        duration_acts = acts[:, :, :, -self.n_durations :]
157
158	        # the - self.sigma here is for logit-undernormalization. Check the paper for details.
159	        label_acts = torch.log_softmax(label_acts, -1) - self.sigma
160
161	        duration_acts = torch.log_softmax(duration_acts, -1)
162
163	        forward_logprob, _ = self.compute_forward_prob(label_acts, duration_acts, labels, act_lens, label_lens)
164	        losses = -forward_logprob
165	        if self.reduction == 'mean_batch':
166	            losses = losses.mean()  # global batch size average
167	        elif self.reduction == 'mean':
168	            losses = torch.div(losses, label_lens).mean()
169	        elif self.reduction == 'sum':
170	            losses = losses.sum()
171	        elif self.reduction == 'mean_volume':
172	            losses = losses.sum() / label_lens.sum()  # same as above but longer samples weigh more
173
174	        return losses
175
176	    def logsumexp(self, a, b):
177	        ret = torch.logsumexp(torch.stack([a, b]), dim=0)
178	        return ret
179
180	    def compute_forward_prob(self, acts, duration_acts, labels, act_lens, label_lens):
181	        """This function implements Equation 7 in the TDT paper https://arxiv.org/pdf/2304.06795.pdf,
182	        Simply put, for each alpha(t, u), it sums over the contribution from all incoming blank arcs and non-blank arcs.
183	        """
184	        B, T, U, _ = acts.shape
185
186	        log_alpha = torch.zeros(B, T, U)
187	        log_alpha = log_alpha.cuda()
188	        for b in range(B):
189	            for t in range(T):
190	                for u in range(U):
191	                    if u == 0:
192	                        if t == 0:
193	                            # both t and u are 0, this is the base case for alphas.
194	                            log_alpha[b, t, u] = 0.0
195	                        else:
196	                            # u = 0 and t != 0: only considers blank emissions.
197	                            log_alpha[b, t, u] = -1000.0
198	                            for n, l in enumerate(self.durations):
199	                                if (
200	                                    t - l >= 0 and l > 0
201	                                ):  # checking conditions for blank emission, l has to be at least 1
202	                                    tmp = (
203	                                        log_alpha[b, t - l, u]
204	                                        + acts[b, t - l, u, self.blank]
205	                                        + duration_acts[b, t - l, u, n]
206	                                    )
207	                                    log_alpha[b, t, u] = self.logsumexp(tmp, 1.0 * log_alpha[b, t, u])
208
209	                    else:
210	                        # u != 0 here, need to consider both blanks and non-blanks.
211	                        log_alpha[b, t, u] = -1000.0
212	                        for n, l in enumerate(self.durations):
213	                            if t - l >= 0:
214	                                if l > 0:  # for blank emissions. Need to ensure index is not out-of-bound.
215	                                    tmp = (
216	                                        log_alpha[b, t - l, u]
217	                                        + acts[b, t - l, u, self.blank]
218	                                        + duration_acts[b, t - l, u, n]
219	                                    )
220	                                    log_alpha[b, t, u] = self.logsumexp(tmp, 1.0 * log_alpha[b, t, u])
221
222	                                # non-blank emissions.
223	                                tmp = (
224	                                    log_alpha[b, t - l, u - 1]
225	                                    + acts[b, t - l, u - 1, labels[b, u - 1]]
226	                                    + duration_acts[b, t - l, u - 1, n]
227	                                )
228	                                log_alpha[b, t, u] = self.logsumexp(tmp, 1.0 * log_alpha[b, t, u])
229
230	        log_probs = []
231	        for b in range(B):
232	            tt = torch.Tensor([-1000.0]).cuda()[0]
233
234	            # need to loop over all possible ways that blank with different durations contributes to the final loss.
235	            for n, l in enumerate(self.durations):
236	                if act_lens[b] - l >= 0 and l > 0:
237	                    bb = (
238	                        log_alpha[b, act_lens[b] - l, label_lens[b]]
239	                        + acts[b, act_lens[b] - l, label_lens[b], self.blank]
240	                        + duration_acts[b, act_lens[b] - l, label_lens[b], n]
241	                    )
242
243	                    tt = self.logsumexp(bb, 1.0 * tt)
244
245	            log_probs.append(tt)
246
247	        log_prob = torch.stack(log_probs)
248
249	        return log_prob, log_alpha
250
251
252	class MultiblankRNNTLossPytorch(Loss):
253	    """
254	    Pure Python implementation of multi-blank transducer loss (https://arxiv.org/pdf/2211.03541.pdf)
255	    """
256
257	    @property
258	    def input_types(self):
259	        """Input types definitions for CTCLoss.
260	        """
261	        return {
262	            "acts": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
263	            "labels": NeuralType(('B', 'T'), LabelsType()),
264	            "act_lens": NeuralType(tuple('B'), LengthsType()),
265	            "label_lens": NeuralType(tuple('B'), LengthsType()),
266	        }
267
268	    @property
269	    def output_types(self):
270	        """Output types definitions for CTCLoss.
271	        loss:
272	            NeuralType(None)
273	        """
274	        return {"loss": NeuralType(elements_type=LossType())}
275
276	    def __init__(self, blank, big_blank_durations, reduction: str = "sum", sigma: float = 0.0):
277	        super().__init__()
278	        self.blank = blank
279	        self.big_blank_durations = big_blank_durations
280	        self.reduction = reduction
281	        self.sigma = sigma
282
283	    def forward(self, acts, labels, act_lens, label_lens):
284	        acts = torch.log_softmax(acts, -1) - self.sigma
285	        forward_logprob, _ = self.compute_forward_prob(acts, labels, act_lens, label_lens)
286
287	        losses = -forward_logprob
288	        if self.reduction == 'mean_batch':
289	            losses = losses.mean()  # global batch size average
290	        elif self.reduction == 'mean':
291	            losses = torch.div(losses, label_lens).mean()
292	        elif self.reduction == 'sum':
293	            losses = losses.sum()
294	        elif self.reduction == 'mean_volume':
295	            losses = losses.sum() / label_lens.sum()  # same as above but longer samples weigh more
296
297	        return losses
298
299	    def compute_forward_prob(self, acts, labels, act_lens, label_lens):
300	        B, T, U, _ = acts.shape
301
302	        log_alpha = torch.zeros(B, T, U, device=acts.device)
303	        for t in range(T):
304	            for u in range(U):
305	                if u == 0:
306	                    if t == 0:
307	                        # this is the base case: (t=0, u=0) with log-alpha = 0.
308	                        log_alpha[:, t, u] = 0.0
309	                    else:
310	                        # this is case for (t = 0, u > 0), reached by (t, u - d)
311	                        # emitting a blank symbol of duration d.
312	                        log_alpha[:, t, u] = log_alpha[:, t - 1, u] + acts[:, t - 1, 0, self.blank]
313	                        for i, d in enumerate(self.big_blank_durations):
314	                            if t >= d:
315	                                tt = log_alpha[:, t - d, u] + acts[:, t - d, 0, self.blank - 1 - i]
316	                                log_alpha[:, t, u] = torch.logsumexp(
317	                                    torch.stack([1.0 * log_alpha[:, t, u], tt]), dim=0
318	                                )
319
320	                else:
321	                    if t == 0:
322	                        # in case of (u > 0, t = 0), this is only reached from
323	                        # (t, u - 1) with a label emission.
324	                        gathered = torch.gather(
325	                            acts[:, t, u - 1], dim=1, index=labels[:, u - 1].view(-1, 1).type(torch.int64)
326	                        ).reshape(-1)
327	                        log_alpha[:, t, u] = log_alpha[:, t, u - 1] + gathered
328	                    else:
329	                        # here both t and u are > 0, this state is reachable
330	                        # with two possibilities: (t - d, u) with emission of
331	                        # blank with duration d, or (t, u - 1) with a label emission.
332
333	                        # first we take care of the standard blank.
334	                        log_alpha[:, t, u] = torch.logsumexp(
335	                            torch.stack(
336	                                [
337	                                    log_alpha[:, t - 1, u] + acts[:, t - 1, u, self.blank],
338	                                    log_alpha[:, t, u - 1]
339	                                    + torch.gather(
340	                                        acts[:, t, u - 1], dim=1, index=labels[:, u - 1].view(-1, 1).type(torch.int64)
341	                                    ).reshape(-1),
342	                                ]
343	                            ),
344	                            dim=0,
345	                        )
346
347	                        # now we go over all big blanks. They need to be considered if current t >= blank duration d.
348	                        for i, d in enumerate(self.big_blank_durations):
349	                            if t >= d:
350	                                tt = log_alpha[:, t - d, u] + acts[:, t - d, u, self.blank - 1 - i]
351	                                log_alpha[:, t, u] = torch.logsumexp(
352	                                    torch.stack([1.0 * log_alpha[:, t, u], tt]), dim=0
353	                                )
354
355	        log_probs = []
356	        for b in range(B):
357	            # here we need to add the final blank emission weights, which needs
358	            # to consider all possible blank durations.
359	            to_append = (
360	                log_alpha[b, act_lens[b] - 1, label_lens[b]] + acts[b, act_lens[b] - 1, label_lens[b], self.blank]
361	            )
362
363	            for i, d in enumerate(self.big_blank_durations):
364	                if act_lens[b] >= d:
365	                    tt = (
366	                        log_alpha[b, act_lens[b] - d, label_lens[b]]
367	                        + acts[b, act_lens[b] - d, label_lens[b], self.blank - 1 - i]
368	                    )
369	                    to_append = torch.logsumexp(torch.stack([1.0 * to_append, tt]), dim=0)
370
371	            log_probs.append(to_append)
372	        log_prob = torch.stack(log_probs)
373
374	        return log_prob, log_alpha
375


1	# ! /usr/bin/python
2	# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	#     http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	import torch
17
18	from nemo.core.classes import Loss, Typing, typecheck
19	from nemo.core.neural_types import LabelsType, LogitsType, LossType, NeuralType
20
21	__all__ = ['AngularSoftmaxLoss']
22
23
24	class AngularSoftmaxLoss(Loss, Typing):
25	    """
26	    Computes ArcFace Angular softmax angle loss
27	    reference: https://openaccess.thecvf.com/content_CVPR_2019/papers/Deng_ArcFace_Additive_Angular_Margin_Loss_for_Deep_Face_Recognition_CVPR_2019_paper.pdf
28	    args:
29	    scale: scale value for cosine angle
30	    margin: margin value added to cosine angle
31	    """
32
33	    @property
34	    def input_types(self):
35	        """Input types definitions for AnguarLoss.
36	        """
37	        return {
38	            "logits": NeuralType(('B', 'D'), LogitsType()),
39	            "labels": NeuralType(('B',), LabelsType()),
40	        }
41
42	    @property
43	    def output_types(self):
44	        """Output types definitions for AngularLoss.
45	        loss:
46	            NeuralType(None)
47	        """
48	        return {"loss": NeuralType(elements_type=LossType())}
49
50	    def __init__(self, scale=20.0, margin=1.35):
51	        super().__init__()
52
53	        self.eps = 1e-7
54	        self.scale = scale
55	        self.margin = margin
56
57	    @typecheck()
58	    def forward(self, logits, labels):
59	        numerator = self.scale * torch.cos(
60	            torch.acos(torch.clamp(torch.diagonal(logits.transpose(0, 1)[labels]), -1.0 + self.eps, 1 - self.eps))
61	            + self.margin
62	        )
63	        excl = torch.cat(
64	            [torch.cat((logits[i, :y], logits[i, y + 1 :])).unsqueeze(0) for i, y in enumerate(labels)], dim=0
65	        )
66	        denominator = torch.exp(numerator) + torch.sum(torch.exp(self.scale * excl), dim=1)
67	        L = numerator - torch.log(denominator)
68	        return -torch.mean(L)
69


Now let me read the remaining loss files and SSL losses:


1	# ! /usr/bin/python
2	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	#     http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	import torch
17
18	from nemo.core.classes import Loss, Typing, typecheck
19	from nemo.core.neural_types import LabelsType, LengthsType, LossType, NeuralType, ProbsType
20
21	__all__ = ['BCELoss']
22
23
24	class BCELoss(Loss, Typing):
25	    """
26	    Computes Binary Cross Entropy (BCE) loss. The BCELoss class expects output from Sigmoid function.
27	    """
28
29	    @property
30	    def input_types(self):
31	        """Input types definitions for AnguarLoss."""
32	        return {
33	            "probs": NeuralType(('B', 'T', 'C'), ProbsType()),
34	            'labels': NeuralType(('B', 'T', 'C'), LabelsType()),
35	            "target_lens": NeuralType(('B'), LengthsType()),
36	        }
37
38	    @property
39	    def output_types(self):
40	        """
41	        Output types definitions for binary cross entropy loss. Weights for labels can be set using weight variables.
42	        """
43	        return {"loss": NeuralType(elements_type=LossType())}
44
45	    def __init__(
46	        self,
47	        reduction: str = 'mean',
48	        alpha: float = 1.0,
49	        weight: torch.Tensor = torch.tensor([0.1, 0.9]),
50	        sorted_preds: bool = False,
51	        sorted_loss: bool = False,
52	        class_normalization: bool = False,
53	    ):
54	        """
55	        A custom loss function that supports class normalization,
56	        weighted binary cross-entropy, and optional sorting.
57
58	        Args:
59	            reduction (str): Specifies the reduction to apply to the output,
60	                options are 'mean', 'sum', or 'none'. Default is 'mean'.
61	            alpha (float): Scaling factor for loss (unused in this implementation). Default is 1.0.
62	            weight (torch.Tensor): Class weights for the binary cross-entropy loss. Default is [0.1, 0.9].
63	            sorted_preds (bool): If True, assumes predictions are sorted. Default is False.
64	            sorted_loss (bool): If True, sorts the loss before reduction. Default is False.
65	            class_normalization (bool): If True, uses 'none' reduction for per-class loss. Default is False.
66	        """
67	        super().__init__()
68	        self.class_normalization = class_normalization
69	        if class_normalization:
70	            self.reduction = 'none'
71	        else:
72	            self.reduction = 'mean'
73	        self.loss_weight = weight
74	        self.loss_f = torch.nn.BCELoss(reduction=self.reduction)
75	        self.sorted_preds = sorted_preds
76	        self.sorted_loss = sorted_loss
77	        self.eps = 1e-6
78
79	    @typecheck()
80	    def forward(self, probs, labels, target_lens):
81	        """
82	        Calculate binary cross entropy loss based on probs, labels and target_lens variables.
83
84	        Args:
85	            probs (torch.tensor)
86	                Predicted probability value which ranges from 0 to 1. Sigmoid output is expected.
87	            labels (torch.tensor)
88	                Groundtruth label for the predicted samples.
89	            target_lens (torch.tensor):
90	                The actual length of the sequence without zero-padding.
91
92	        Returns:
93	            loss (NeuralType)
94	                Binary cross entropy loss value.
95	        """
96	        probs_list = [probs[k, : target_lens[k], :] for k in range(probs.shape[0])]
97	        targets_list = [labels[k, : target_lens[k], :] for k in range(labels.shape[0])]
98	        probs = torch.cat(probs_list, dim=0)
99	        labels = torch.cat(targets_list, dim=0)
100	        norm_weight = torch.zeros_like(labels).detach().clone()
101	        loss = torch.tensor(0.0).to(labels.device)
102
103	        if self.class_normalization in ['class', 'class_binary', 'binary']:
104	            if self.class_normalization in ['class', 'class_binary']:
105	                # Normalize loss by number of classes
106	                norm_weight = 1 / (labels.sum(dim=0) + self.eps)
107	                norm_weight_norm = norm_weight / norm_weight.sum()
108	                norm_weight_norm = torch.clamp(norm_weight_norm, min=0.05, max=1.0)
109	                norm_weight_norm = norm_weight_norm / norm_weight_norm.max()
110	                norm_weight = norm_weight_norm[None, :].expand_as(labels).detach().clone()
111	            else:
112	                norm_weight = torch.ones_like(labels).detach().clone()
113
114	            if self.class_normalization in ['binary', 'class_binary']:
115	                binary_weight = torch.ones_like(labels).detach().clone()
116	                one_weight = (labels.sum() / (labels.shape[0] * labels.shape[1])).to(labels.device)
117	                binary_weight[labels == 0] = one_weight
118	                binary_weight[labels == 1] = 1 - one_weight
119	            else:
120	                binary_weight = torch.ones_like(labels).detach().clone()
121
122	        elif self.class_normalization == 'none' or not self.class_normalization:
123	            binary_weight = torch.ones_like(labels).detach().clone()
124	            norm_weight = torch.ones_like(labels).detach().clone()
125
126	        if self.reduction == 'sum':
127	            loss = self.loss_f(probs, labels)
128	        elif self.reduction == 'mean':
129	            loss = self.loss_f(probs, labels).mean()
130	        elif self.reduction == 'none':
131	            if self.class_normalization in ['class', 'class_binary', 'binary']:
132	                loss = (binary_weight * norm_weight * self.loss_f(probs, labels)).sum()
133	            else:
134	                loss = self.loss_f(probs, labels)
135	        return loss
136


1	# ! /usr/bin/python
2	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	#     http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	from typing import Optional
17
18	import torch
19	from omegaconf import DictConfig
20
21	from nemo.core.classes import Loss, typecheck
22	from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
23
24
25	class LatticeLoss(Loss):
26	    """Family of loss functions based on various lattice scores.
27
28	    Note:
29	        Requires k2 v1.14 or later to be installed to use this loss function.
30
31	    Losses can be selected via the config, and optionally be passed keyword arguments as follows.
32
33	    Examples:
34	        .. code-block:: yaml
35
36	            model:  # Model config
37	                ...
38	                graph_module_cfg:  # Config for graph modules, e.g. LatticeLoss
39	                    criterion_type: "map"
40	                    loss_type: "mmi"
41	                    split_batch_size: 0
42	                    backend_cfg:
43	                        topo_type: "default"       # other options: "compact", "shared_blank", "minimal"
44	                        topo_with_self_loops: true
45	                        token_lm: <token_lm_path>  # must be provided for criterion_type: "map"
46
47	    Args:
48	        num_classes: Number of target classes for the decoder network to predict.
49	            (Excluding the blank token).
50
51	        reduction: Type of reduction to perform on loss. Possible values are `mean_batch`, `mean`, `sum`, or None.
52	            None will return a torch vector comprising the individual loss values of the batch.
53
54	        backend: Which backend to use for loss calculation. Currently only `k2` is supported.
55
56	        criterion_type: Type of criterion to use. Choices: `ml` and `map`,
57	            with `ml` standing for Maximum Likelihood and `map` for Maximum A Posteriori Probability.
58
59	        loss_type: Type of the loss function to use. Choices: `ctc` and `rnnt` for `ml`, and `mmi` for `map`.
60
61	        split_batch_size: Local batch size. Used for memory consumption reduction at the cost of speed performance.
62	            Effective if complies 0 < split_batch_size < batch_size.
63
64	        graph_module_cfg: Optional Dict of (str, value) pairs that are passed to the backend loss function.
65	    """
66
67	    @property
68	    def input_types(self):
69	        """Input types definitions for LatticeLoss.
70	        """
71	        return {
72	            "log_probs": NeuralType(("B", "T", "D") if self._3d_input else ("B", "T", "T", "D"), LogprobsType()),
73	            "targets": NeuralType(("B", "T"), LabelsType()),
74	            "input_lengths": NeuralType(tuple("B"), LengthsType()),
75	            "target_lengths": NeuralType(tuple("B"), LengthsType()),
76	        }
77
78	    @property
79	    def output_types(self):
80	        """Output types definitions for LatticeLoss.
81	        loss:
82	            NeuralType(None)
83	        """
84	        return {"loss": NeuralType(elements_type=LossType())}
85
86	    def __init__(
87	        self,
88	        num_classes: int,
89	        reduction: str = "mean_batch",
90	        backend: str = "k2",
91	        criterion_type: str = "ml",
92	        loss_type: str = "ctc",
93	        split_batch_size: int = 0,
94	        graph_module_cfg: Optional[DictConfig] = None,
95	    ):
96	        super().__init__()
97	        self._blank = num_classes
98	        self.split_batch_size = split_batch_size
99	        inner_reduction = None
100	        if reduction == "mean_batch":
101	            inner_reduction = "none"
102	            self._apply_batch_mean = True
103	        elif reduction in ["sum", "mean", "none"]:
104	            inner_reduction = reduction
105	            self._apply_batch_mean = False
106
107	        # we assume that self._blank + 1 == num_classes
108	        if backend == "k2":
109	            if criterion_type == "ml":
110	                if loss_type == "ctc":
111	                    from nemo.collections.asr.parts.k2.ml_loss import CtcLoss as K2Loss
112	                elif loss_type == "rnnt":
113	                    from nemo.collections.asr.parts.k2.ml_loss import RnntLoss as K2Loss
114	                else:
115	                    raise ValueError(f"Unsupported `loss_type`: {loss_type}.")
116	            elif criterion_type == "map":
117	                if loss_type == "ctc":
118	                    from nemo.collections.asr.parts.k2.map_loss import CtcMmiLoss as K2Loss
119	                else:
120	                    raise ValueError(f"Unsupported `loss_type`: {loss_type}.")
121	            else:
122	                raise ValueError(f"Unsupported `criterion_type`: {criterion_type}.")
123
124	            self._loss = K2Loss(
125	                num_classes=self._blank + 1, blank=self._blank, reduction=inner_reduction, cfg=graph_module_cfg,
126	            )
127	        elif backend == "gtn":
128	            raise NotImplementedError(f"Backend {backend} is not supported.")
129	        else:
130	            raise ValueError(f"Invalid value of `backend`: {backend}.")
131
132	        self.criterion_type = criterion_type
133	        self.loss_type = loss_type
134	        self._3d_input = self.loss_type != "rnnt"
135
136	        if self.split_batch_size > 0:
137	            # don't need to guard grad_utils
138	            from nemo.collections.asr.parts.k2.grad_utils import PartialGrad
139
140	            self._partial_loss = PartialGrad(self._loss)
141
142	    def update_graph(self, graph):
143	        """Updates graph of the backend loss function.
144	        """
145	        if self.criterion_type != "ml":
146	            self._loss.update_graph(graph)
147
148	    @typecheck()
149	    def forward(self, log_probs, targets, input_lengths, target_lengths):
150	        # override forward implementation
151	        # custom logic, if necessary
152
153	        assert not (torch.isnan(log_probs).any() or torch.isinf(log_probs).any())
154
155	        log_probs = log_probs.float()
156	        input_lengths = input_lengths.long()
157	        target_lengths = target_lengths.long()
158	        targets = targets.long()
159	        batch_size = log_probs.shape[0]
160	        if self.split_batch_size > 0 and self.split_batch_size <= batch_size:
161	            loss_list = []
162	            for batch_idx in range(0, batch_size, self.split_batch_size):
163	                begin = batch_idx
164	                end = min(begin + self.split_batch_size, batch_size)
165	                input_lengths_part = input_lengths[begin:end]
166	                log_probs_part = log_probs[begin:end, : input_lengths_part.max()]
167	                target_lengths_part = target_lengths[begin:end]
168	                targets_part = targets[begin:end, : target_lengths_part.max()]
169	                loss_part, _ = (
170	                    self._partial_loss(log_probs_part, targets_part, input_lengths_part, target_lengths_part)
171	                    if log_probs_part.requires_grad
172	                    else self._loss(log_probs_part, targets_part, input_lengths_part, target_lengths_part)
173	                )
174	                del log_probs_part, targets_part, input_lengths_part, target_lengths_part
175	                loss_list.append(loss_part)
176	            loss = torch.cat(loss_list, 0)
177	        else:
178	            loss, _ = self._loss(
179	                log_probs=log_probs, targets=targets, input_lengths=input_lengths, target_lengths=target_lengths,
180	            )
181	        if self._apply_batch_mean:
182	            # torch.mean gives nan if loss is empty
183	            loss = torch.mean(loss) if loss.nelement() > 0 else torch.sum(loss)
184	        return loss
185


1	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	from math import ceil
16
17	import torch
18	import torch.nn.functional as F
19	from torch import nn
20
21	from nemo.core import Loss, typecheck
22	from nemo.core.neural_types import AcousticEncodedRepresentation, LengthsType, LossType, NeuralType, SpectrogramType
23
24	__all__ = ["ContrastiveLoss"]
25
26
27	class ContrastiveLoss(Loss):
28	    @property
29	    def input_types(self):
30	        """Input types definitions for Contrastive."""
31	        return {
32	            "spectrograms": NeuralType(("B", "D", "T"), SpectrogramType()),
33	            "spec_masks": NeuralType(("B", "D", "T"), SpectrogramType()),
34	            "decoder_outputs": NeuralType(("B", "T", "D"), AcousticEncodedRepresentation()),
35	            "decoder_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
36	        }
37
38	    @property
39	    def output_types(self):
40	        """Output types definitions for Contrastive.
41	        loss:
42	            NeuralType(None)
43	        """
44	        return {"loss": NeuralType(elements_type=LossType())}
45
46	    @property
47	    def needs_labels(self):
48	        return False
49
50	    def __init__(
51	        self,
52	        in_dim: int,
53	        proj_dim: int = 128,
54	        combine_time_steps: int = 1,
55	        num_negatives: int = 100,
56	        quantized_targets: bool = False,
57	        codebook_size: int = 320,
58	        prob_ppl_weight: float = 0.1,
59	        logit_temp: float = 0.1,
60	        reduce: str = "sum",
61	        sample_from_same_utterance_only: bool = True,
62	        sample_from_non_masked: bool = False,
63	        sample_from_codebook: bool = False,
64	        group_loss: bool = False,
65	        num_groups: int = 2,
66	        quantizer_temp_start: float = 2,
67	        quantizer_temp_min: float = 0.5,
68	        quantizer_temp_decay: float = 0.999995,
69	        mask_threshold: float = 0.8,
70	        store_ids: bool = True,
71	        reduce_ids: bool = False,
72	        multiplier: float = 16.0,
73	    ):
74	        """
75	        Loss function representing the contrastive task of identifying the true latent speech representation of
76	        the masked spectrogram steps from a set of sampled distractors.
77
78	        Args:
79	            in_dim: Number of spectrogram channels.
80	            proj_dim: Number of channels in the model outputs.
81	            combine_time_steps: How many time steps should be combined into a single representation.
82	            num_negatives: Number of sampled negatives for each target.
83	            quantized_targets: Bool that determines if the targets should be quantized.
84	            codebook_size: Number of vectors in the codebook per group.
85	            prob_ppl_weight: Float multiplier on the perplexity loss for target quantization.
86	            logit_temp: Float temperature for normalizing logits.
87	            reduce: String representing the type of reduction used for cross entropy.
88	            sample_from_same_utterance_only: Bool that determines if negatives should be sampled only from same utterance.
89	            sample_from_non_masked: Bool that determines if negatives should be sampled from non-masked steps of the spectrogram.
90	            sample_from_codebook: Bool that determines if negatives should be sampled from entire codebook.
91	            group_loss: Bool that determines if loss should be computed separately for each group in the quantizer codebook.
92	            num_groups: Number of groups in the quantizer codebook.
93	            quantizer_temp_start: Starting temperature in quantizer.
94	            quantizer_temp_min: Minimum temperature in quantizer.
95	            quantizer_temp_decay: Decay rate of quantizer temperature per global step.
96	            mask_threshold: Float threshold for determining if a time step of the spectrogram is masked based on percent of masked channels.
97	            store_ids: Bool that determines if the quantizer ids will be stored to be potentially used by other losses.
98	            reduce_ids: Bool that determines if we convert any sequence of consecutive equivalent ids to a single occurence of that id.
99	            multiplier: Float multipler on final loss
100	        """
101
102	        super().__init__()
103	        quantizer_temp = (quantizer_temp_start, quantizer_temp_min, quantizer_temp_decay)
104	        self.quantized_targets = quantized_targets
105	        self.num_negatives = num_negatives
106	        self.prob_ppl_weight = prob_ppl_weight
107	        if self.quantized_targets:
108	            quantizer_cfg = {
109	                "_target_": "nemo.collections.asr.parts.submodules.ssl_quantizers.GumbelVectorQuantizer",
110	                "dim": in_dim * combine_time_steps,
111	                "vq_dim": proj_dim,
112	                "num_vars": codebook_size,
113	                "groups": num_groups,
114	                "temp": quantizer_temp,
115	                "combine_groups": True,
116	                "time_first": True,
117	            }
118	            self.quantizer = ContrastiveLoss.from_config_dict(quantizer_cfg)
119	        self.prob_ppl_weight = prob_ppl_weight
120	        self.logit_temp = logit_temp
121	        self.reduce = reduce
122	        self.combine_time_steps = combine_time_steps
123	        self.sample_from_same_utterance_only = sample_from_same_utterance_only
124	        self.sample_from_non_masked = sample_from_non_masked
125	        self.sample_from_codebook = sample_from_codebook
126	        self.group_loss = group_loss
127	        self.mask_threshold = mask_threshold
128	        self.multiplier = multiplier
129
130	        self.store_ids = store_ids
131	        self.reduce_ids = reduce_ids
132
133	        if not self.quantized_targets:
134	            self.target_proj = nn.Linear(in_dim * combine_time_steps, proj_dim)
135
136	    def sample_negatives(self, y, num):
137	        # y - T'xBxC or T'xC
138
139	        high = y.shape[0]
140	        neg_idxs = torch.multinomial(torch.ones((num, high), device=y.device), self.num_negatives)
141
142	        negs = y[neg_idxs.view(-1)]
143	        negs = negs.view((num, self.num_negatives) + y.shape[1:])
144	        negs = negs.transpose(0, 1)
145	        # negs - NxT'xBxC or NxT'xC
146
147	        return negs, neg_idxs
148
149	    @typecheck()
150	    def forward(self, spectrograms, spec_masks, decoder_outputs, decoder_lengths=None):
151	        targets = spectrograms.transpose(-2, -1)
152	        masks = spec_masks.transpose(-2, -1)
153	        # BxTxC
154	        diff = int(ceil(targets.shape[1] / decoder_outputs.shape[1]) * decoder_outputs.shape[1]) - targets.shape[1]
155
156	        if diff > 0:
157	            targets = F.pad(targets, (0, 0, 0, diff))
158	            masks = F.pad(masks, (0, 0, 0, diff))
159
160	        targets = targets.reshape(targets.shape[0], decoder_outputs.shape[1], -1)
161	        masks = masks.reshape(targets.shape[0], decoder_outputs.shape[1], -1)
162
163	        if self.quantized_targets:
164	            if self.store_ids:
165	                # store ids for use by other losses
166	                targets, prob_ppl_loss, cur_codebook_temp, self.target_ids = self.quantizer(targets, return_ids=True)
167
168	                if self.reduce_ids:
169	                    # reduce consecutive equivalent ids to a single occurence
170	                    _, indices = torch.unique_consecutive(self.target_ids, return_inverse=True)
171	                    indices -= indices.min(dim=1, keepdims=True)[0]
172	                    reduced_ids = torch.zeros_like(self.target_ids)
173	                    reduced_ids = reduced_ids.scatter_(1, indices, self.target_ids)
174	                    reduced_lens = indices.max(dim=-1)[0] + 1
175
176	                    self.target_ids = reduced_ids.narrow(1, 0, reduced_lens.max())
177	                    self.target_lengths = reduced_lens
178
179	                else:
180	                    self.target_lengths = None
181
182	            else:
183	                targets, prob_ppl_loss, cur_codebook_temp = self.quantizer(targets)
184	        else:
185	            targets = self.target_proj(targets)
186
187	        if self.sample_from_same_utterance_only:
188	            bs = decoder_outputs.shape[0]
189	            masks = masks.mean(-1) > self.mask_threshold
190	            out_masked_only = decoder_outputs[masks]
191	            targets_masked_only = targets[masks]
192	            out_masked_only = out_masked_only.reshape(bs, -1, out_masked_only.shape[-1])
193	            targets_masked_only = targets_masked_only.reshape(bs, -1, targets_masked_only.shape[-1])
194
195	            # BxT'xC
196	            # number of masked time steps to predict (T')
197	            # -> T'xBxC
198
199	            out_masked_only = out_masked_only.transpose(0, 1)
200	            targets_masked_only = targets_masked_only.transpose(0, 1)
201	            # -> T'xBxC
202
203	            if self.sample_from_non_masked:
204	                # sample from all steps in utterance
205	                negatives, _ = self.sample_negatives(
206	                    targets.transpose(0, 1),
207	                    targets_masked_only.size(0),  # TxBxC  # T'
208	                )
209	            else:
210	                # only sample from masked steps in utterance
211	                negatives, _ = self.sample_negatives(targets_masked_only, targets_masked_only.size(0))  # T'xBxC  # T'
212	            # NxT'xBxC
213
214	            out_masked_only = out_masked_only.reshape(-1, out_masked_only.shape[-1])
215	            targets_masked_only = targets_masked_only.reshape(-1, targets_masked_only.shape[-1])
216	            negatives = negatives.reshape(self.num_negatives, -1, negatives.shape[-1])
217
218	            # T'BxC and NxT'BxC
219
220	        else:
221	            masks = masks.mean(-1) > self.mask_threshold
222	            out_masked_only = decoder_outputs[masks]
223	            targets_masked_only = targets[masks]
224
225	            # T'xC
226	            # number of masked time steps to predict (T')
227
228	            if self.group_loss:
229	                num_groups = self.quantizer.groups
230	                negatives = self.quantizer.vars.reshape(num_groups, self.quantizer.num_vars, -1)
231	                # GxNx(C//G)
232	                negatives = negatives.transpose(0, 1)
233	                # NxGx(C//G)
234	                negatives = negatives.unsqueeze(1).expand(-1, out_masked_only.shape[0], -1, -1)
235	                # NxT'xGx(C//G)
236	                negatives = negatives.reshape(negatives.shape[0], -1, negatives.shape[-1])
237	                # NxT'Gx(C//G)
238
239	                out_masked_only = out_masked_only.reshape(-1, out_masked_only.shape[-1] // num_groups)
240	                targets_masked_only = targets_masked_only.reshape(-1, targets_masked_only.shape[-1] // num_groups)
241	                # T'Gx(C//G)
242	            elif self.sample_from_codebook:
243	                # sample from the full codebook
244	                negatives = self.quantizer.sample_from_codebook(self.num_negatives, targets_masked_only.size(0))
245	            elif self.sample_from_non_masked:
246	                # sample from all steps in batch
247	                negatives, _ = self.sample_negatives(
248	                    targets.reshape(targets.shape[0] * targets.shape[1], -1),
249	                    targets_masked_only.size(0),  # BTxC
250	                )  # T'
251	            else:
252	                # only sample from masked steps
253	                negatives, _ = self.sample_negatives(targets_masked_only, targets_masked_only.size(0))  # T'xC  # T'
254	                # NxT'xC
255
256	        # Calculate similarity between outputs and all targets
257	        similarity_scores = self._calculate_similarity(out_masked_only, negatives, targets_masked_only)
258	        # (1+N)xT'
259	        # cosine similarity of outs with targets + N negatives
260
261	        # Create targets of size T
262	        similarity_targets = decoder_outputs.new_zeros(similarity_scores.size(1), dtype=torch.long)
263	        # T'
264	        # targets are 0, since it's the first, followed by N sampled negatives
265
266	        # Transpose similarity scores to TxF for loss
267	        similarity_scores = similarity_scores.transpose(0, 1)
268	        # T'x(1+N)
269
270	        loss = F.cross_entropy(similarity_scores, similarity_targets, reduction=self.reduce)
271
272	        sample_size = similarity_targets.numel()
273
274	        if self.prob_ppl_weight != 0 and self.quantized_targets:
275	            prob_ppl_loss = self.prob_ppl_weight * prob_ppl_loss * sample_size
276	            loss += prob_ppl_loss
277
278	        if not isinstance(loss, torch.Tensor):
279	            loss = torch.Tensor([0]).to(device=decoder_outputs.device)
280
281	        batch_size = spectrograms.shape[0]
282	        loss *= self.multiplier / batch_size
283
284	        return loss
285
286	    def _calculate_similarity(self, logits, negatives, targets):
287	        neg_is_pos = (targets == negatives).all(-1)
288	        # NxT' - true where the negative is actually the positive
289	        targets = targets.unsqueeze(0)
290	        # 1xT'xC
291	        targets = torch.cat([targets, negatives], dim=0)
292	        # (1+N)xT'XC
293	        logits = torch.cosine_similarity(
294	            logits.float().unsqueeze(0).expand(targets.shape[0], -1, -1), targets.float(), dim=-1
295	        ).type_as(logits)
296	        # (1+N)xT'
297	        logits /= self.logit_temp
298	        if neg_is_pos.any():
299	            logits[1:][neg_is_pos] = float("-inf")
300	        return logits
301
302	    def set_num_updates(self, num_updates):
303	        if self.quantized_targets:
304	            self.quantizer.set_num_updates(num_updates)
305


Now let me check the SSL losses:


1	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	from nemo.collections.asr.losses import CTCLoss
16	from nemo.core import Loss, typecheck
17	from nemo.core.neural_types import LabelsType, LengthsType, LossType, NeuralType, SpectrogramType, VoidType
18
19	__all__ = ["CTCLossForSSL"]
20
21
22	class CTCLossForSSL(Loss):
23	    @property
24	    def input_types(self):
25	        """Input types definitions for Contrastive.
26	        """
27	        return {
28	            "spec_masks": NeuralType(("B", "D", "T"), SpectrogramType()),
29	            "decoder_outputs": NeuralType(("B", "T", "D"), VoidType()),
30	            "targets": NeuralType(('B', 'T'), LabelsType()),
31	            "decoder_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
32	            "target_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
33	        }
34
35	    @property
36	    def output_types(self):
37	        """Output types definitions for Contrastive.
38	        loss:
39	            NeuralType(None)
40	        """
41	        return {"loss": NeuralType(elements_type=LossType())}
42
43	    @property
44	    def needs_labels(self):
45	        return True
46
47	    def __init__(self, num_classes, zero_infinity=True, reduction='mean_batch'):
48	        super().__init__()
49	        self.loss = CTCLoss(num_classes=num_classes, reduction=reduction, zero_infinity=zero_infinity)
50
51	    @typecheck()
52	    def forward(self, spec_masks, decoder_outputs, targets, decoder_lengths=None, target_lengths=None):
53	        loss = self.loss(
54	            log_probs=decoder_outputs, targets=targets, input_lengths=decoder_lengths, target_lengths=target_lengths
55	        )
56
57	        return loss
58


1	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	import torch
16	import torch.nn.functional as F
17	from torch import nn
18
19	from nemo.core import Loss, typecheck
20	from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType, SpectrogramType
21
22	__all__ = ["MLMLoss"]
23
24
25	class MLMLoss(Loss):
26	    @property
27	    def input_types(self):
28	        """Input types definitions for Contrastive."""
29	        return {
30	            "spec_masks": NeuralType(("B", "D", "T"), SpectrogramType(), optional=True),
31	            "decoder_outputs": NeuralType(("B", "T", "D"), LogprobsType()),
32	            "targets": NeuralType(('B', 'T'), LabelsType()),
33	            "decoder_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
34	            "target_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
35	            "masks": NeuralType(("B", "D", "T"), SpectrogramType(), optional=True),
36	        }
37
38	    @property
39	    def output_types(self):
40	        """Output types definitions for Contrastive.
41	        loss:
42	            NeuralType(None)
43	        """
44	        return {"loss": NeuralType(elements_type=LossType())}
45
46	    @property
47	    def needs_labels(self):
48	        return True
49
50	    def __init__(
51	        self,
52	        combine_time_steps: int = 1,
53	        mask_threshold: float = 0.8,
54	    ):
55	        super().__init__()
56	        self.nll_loss = nn.NLLLoss()
57	        self.combine_time_steps = combine_time_steps
58	        self.mask_threshold = mask_threshold
59
60	    @typecheck()
61	    def forward(
62	        self, decoder_outputs, targets, decoder_lengths=None, target_lengths=None, spec_masks=None, masks=None
63	    ):
64
65	        if masks is None:
66	            masks = spec_masks
67
68	        # B,D,T -> B,T,D
69	        masks = masks.transpose(1, 2)
70
71	        masks = masks.reshape(masks.shape[0], masks.shape[1] // self.combine_time_steps, -1)
72	        masks = masks.mean(-1) > self.mask_threshold
73
74	        out_masked_only = decoder_outputs[masks]
75	        targets = F.pad(targets, (0, masks.shape[-1] - targets.shape[-1]))
76	        targets_masked_only = targets[masks]
77
78	        loss = self.nll_loss(out_masked_only, targets_masked_only)
79	        loss = torch.mean(loss)
80
81	        return loss
82
83
84	class MultiMLMLoss(Loss):
85	    """
86	    Masked language model loss for multiple decoders, where cross-entropy loss is applied separately on each decoder.
87	    This loss can be used with `nemo.collections.asr.modules.ssl_modules.MultiSoftmaxDecoder` to train a model with multiple targets per frame.
88	    Reference: https://arxiv.org/abs/2202.01855
89	    """
90
91	    @property
92	    def input_types(self):
93	        if self.squeeze_single and self.num_decoders == 1:
94	            decoder_outputs = NeuralType(("B", "T", "C"), LogprobsType())
95	            targets = NeuralType(('B', 'T'), LabelsType())
96	        else:
97	            decoder_outputs = NeuralType(("B", "T", "C", "H"), LogprobsType())
98	            targets = NeuralType(("B", "T", "H"), LabelsType())
99	        return {
100	            "masks": NeuralType(("B", "D", "T"), SpectrogramType()),
101	            "decoder_outputs": decoder_outputs,
102	            "targets": targets,
103	            "decoder_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
104	            "target_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
105	        }
106
107	    def __init__(
108	        self,
109	        combine_time_steps: int = 1,
110	        mask_threshold: float = 0.8,
111	        num_decoders: int = 1,
112	        squeeze_single: bool = False,
113	    ):
114	        super().__init__()
115	        self.num_decoders = num_decoders
116	        self.squeeze_single = squeeze_single
117	        self.mlm_loss = MLMLoss(combine_time_steps, mask_threshold)
118
119	    @typecheck()
120	    def forward(self, masks, decoder_outputs, targets, decoder_lengths=None, target_lengths=None):
121	        if self.squeeze_single and self.num_decoders == 1:
122	            return self.mlm_loss(
123	                spec_masks=masks,
124	                decoder_outputs=decoder_outputs,
125	                targets=targets,
126	                decoder_lengths=decoder_lengths,
127	                target_lengths=target_lengths,
128	            )
129	        loss = 0.0
130	        for i in range(self.num_decoders):
131	            loss += self.mlm_loss(
132	                spec_masks=masks,
133	                decoder_outputs=decoder_outputs[:, :, :, i],
134	                targets=targets[:, :, i],
135	                decoder_lengths=decoder_lengths,
136	                target_lengths=target_lengths,
137	            )
138	        return loss / self.num_decoders
139


1	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	from nemo.collections.asr.losses.rnnt import RNNTLoss
16	from nemo.core import Loss, typecheck
17	from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType, SpectrogramType
18
19	__all__ = ["RNNTLossForSSL"]
20
21
22	class RNNTLossForSSL(Loss):
23	    @property
24	    def input_types(self):
25	        """Input types definitions for Contrastive.
26	        """
27	        return {
28	            "spec_masks": NeuralType(("B", "D", "T"), SpectrogramType()),
29	            "decoder_outputs": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
30	            "targets": NeuralType(('B', 'T'), LabelsType()),
31	            "decoder_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
32	            "target_lengths": NeuralType(tuple('B'), LengthsType(), optional=True),
33	        }
34
35	    @property
36	    def output_types(self):
37	        """Output types definitions for Contrastive.
38	        loss:
39	            NeuralType(None)
40	        """
41	        return {"loss": NeuralType(elements_type=LossType())}
42
43	    @property
44	    def needs_labels(self):
45	        return True
46
47	    def __init__(self, num_classes):
48	        super().__init__()
49	        self.loss = RNNTLoss(num_classes=num_classes)
50
51	    @typecheck()
52	    def forward(self, spec_masks, decoder_outputs, targets, decoder_lengths=None, target_lengths=None):
53
54	        loss = self.loss(
55	            log_probs=decoder_outputs, targets=targets, input_lengths=decoder_lengths, target_lengths=target_lengths
56	        )
57
58	        return loss
59


Now let me check the hybrid models to see how they use losses:


1	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	import copy
16	import os
17	from typing import Dict, List, Optional, Union
18
19	import torch
20	from lightning.pytorch import Trainer
21	from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict
22
23	from nemo.collections.asr.data import audio_to_text_dataset
24	from nemo.collections.asr.data.audio_to_text import _AudioTextDataset
25	from nemo.collections.asr.data.audio_to_text_dali import AudioToBPEDALIDataset
26	from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset
27	from nemo.collections.asr.losses.ctc import CTCLoss
28	from nemo.collections.asr.losses.rnnt import RNNTLoss
29	from nemo.collections.asr.metrics.wer import WER
30	from nemo.collections.asr.models.hybrid_rnnt_ctc_models import EncDecHybridRNNTCTCModel
31	from nemo.collections.asr.parts.mixins import ASRBPEMixin
32	from nemo.collections.asr.parts.submodules.ctc_decoding import CTCBPEDecoding, CTCBPEDecodingConfig
33	from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTBPEDecoding, RNNTBPEDecodingConfig
34	from nemo.collections.asr.parts.utils.asr_batching import get_semi_sorted_batch_sampler
35	from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
36	from nemo.core.classes.common import PretrainedModelInfo
37	from nemo.utils import logging, model_utils
38
39
40	class EncDecHybridRNNTCTCBPEModel(EncDecHybridRNNTCTCModel, ASRBPEMixin):
41	    """Base class for encoder decoder RNNT-based models with auxiliary CTC decoder/loss and subword tokenization."""
42
43	    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
44	        # Convert to Hydra 1.0 compatible DictConfig
45	        cfg = model_utils.convert_model_config_to_dict_config(cfg)
46	        cfg = model_utils.maybe_update_config_version(cfg)
47
48	        # Tokenizer is necessary for this model
49	        if 'tokenizer' not in cfg:
50	            raise ValueError("`cfg` must have `tokenizer` config to create a tokenizer !")
51
52	        if not isinstance(cfg, DictConfig):
53	            cfg = OmegaConf.create(cfg)
54
55	        # Setup the tokenizer
56	        self._setup_tokenizer(cfg.tokenizer)
57
58	        # Initialize a dummy vocabulary
59	        vocabulary = self.tokenizer.tokenizer.get_vocab()
60
61	        # Set the new vocabulary
62	        with open_dict(cfg):
63	            cfg.labels = ListConfig(list(vocabulary))
64
65	        with open_dict(cfg.decoder):
66	            cfg.decoder.vocab_size = len(vocabulary)
67
68	        with open_dict(cfg.joint):
69	            cfg.joint.num_classes = len(vocabulary)
70	            cfg.joint.vocabulary = ListConfig(list(vocabulary))
71	            cfg.joint.jointnet.encoder_hidden = cfg.model_defaults.enc_hidden
72	            cfg.joint.jointnet.pred_hidden = cfg.model_defaults.pred_hidden
73
74	        # setup auxiliary CTC decoder
75	        if 'aux_ctc' not in cfg:
76	            raise ValueError(
77	                "The config need to have a section for the CTC decoder named as aux_ctc for Hybrid models."
78	            )
79
80	        with open_dict(cfg):
81	            if self.tokenizer_type == "agg":
82	                cfg.aux_ctc.decoder.vocabulary = ListConfig(vocabulary)
83	            else:
84	                cfg.aux_ctc.decoder.vocabulary = ListConfig(list(vocabulary.keys()))
85
86	        if cfg.aux_ctc.decoder["num_classes"] < 1:
87	            logging.info(
88	                "\nReplacing placholder number of classes ({}) with actual number of classes - {}".format(
89	                    cfg.aux_ctc.decoder["num_classes"], len(vocabulary)
90	                )
91	            )
92	            cfg.aux_ctc.decoder["num_classes"] = len(vocabulary)
93
94	        super().__init__(cfg=cfg, trainer=trainer)
95
96	        self.cfg.decoding = self.set_decoding_type_according_to_loss(self.cfg.decoding)
97	        # Setup decoding object
98	        self.decoding = RNNTBPEDecoding(
99	            decoding_cfg=self.cfg.decoding,
100	            decoder=self.decoder,
101	            joint=self.joint,
102	            tokenizer=self.tokenizer,
103	        )
104
105	        # Setup wer object
106	        self.wer = WER(
107	            decoding=self.decoding,
108	            batch_dim_index=0,
109	            use_cer=self.cfg.get('use_cer', False),
110	            log_prediction=self.cfg.get('log_prediction', True),
111	            dist_sync_on_step=True,
112	        )
113
114	        # Setup fused Joint step if flag is set
115	        if self.joint.fuse_loss_wer:
116	            self.joint.set_loss(self.loss)
117	            self.joint.set_wer(self.wer)
118
119	        # Setup CTC decoding
120	        ctc_decoding_cfg = self.cfg.aux_ctc.get('decoding', None)
121	        if ctc_decoding_cfg is None:
122	            ctc_decoding_cfg = OmegaConf.structured(CTCBPEDecodingConfig)
123	            with open_dict(self.cfg.aux_ctc):
124	                self.cfg.aux_ctc.decoding = ctc_decoding_cfg
125	        self.ctc_decoding = CTCBPEDecoding(self.cfg.aux_ctc.decoding, tokenizer=self.tokenizer)
126
127	        # Setup CTC WER
128	        self.ctc_wer = WER(
129	            decoding=self.ctc_decoding,
130	            use_cer=self.cfg.aux_ctc.get('use_cer', False),
131	            dist_sync_on_step=True,
132	            log_prediction=self.cfg.get("log_prediction", False),
133	        )
134
135	        # setting the RNNT decoder as the default one
136	        self.cur_decoder = "rnnt"
137
138	    def _setup_dataloader_from_config(self, config: Optional[Dict]):
139
140	        if config.get("use_lhotse"):
141	            return get_lhotse_dataloader_from_config(
142	                config,
143	                # During transcription, the model is initially loaded on the CPU.
144	                # To ensure the correct global_rank and world_size are set,
145	                # these values must be passed from the configuration.
146	                global_rank=self.global_rank if not config.get("do_transcribe", False) else config.get("global_rank"),
147	                world_size=self.world_size if not config.get("do_transcribe", False) else config.get("world_size"),
148	                dataset=LhotseSpeechToTextBpeDataset(
149	                    tokenizer=self.tokenizer,
150	                    return_cuts=config.get("do_transcribe", False),


grep -n "auxiliary\|multi.*loss\|auxiliary.*loss" /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/hybrid_rnnt_ctc_bpe_models.py | head -20


41:    """Base class for encoder decoder RNNT-based models with auxiliary CTC decoder/loss and subword tokenization."""
74:        # setup auxiliary CTC decoder
271:            ctc_decoding_cfg: A config for auxiliary CTC decoding, which is optional and can be used to change the decoding type.


Let me check for inter-CTC and mixin files: