Pool stats with padding
State: closed | Pull request: yes
Labels: name: ASR

# What does this PR do ?

Fix for #5401 

Compute padded stats in pooling layer.

When `lengths` is `None`, the old behavior is preserved.

**Collection**: ASR

# Usage

```python
from typing import List

import torch
import torch.nn as nn
from pytorch_lightning import seed_everything

from nemo.collections.asr.parts.submodules.tdnn_attention import StatsPoolLayer


class OldStatsPoolLayer(nn.Module):
    """Preserved copy of nemo.collections.asr.parts.submodules.tdnn_attention.StatsPoolLayer"""
    def __init__(self, feat_in: int, pool_mode: str = 'xvector'):
        super().__init__()
        self.pool_mode = pool_mode
        self.feat_in = feat_in
        if self.pool_mode == 'xvector':
            self.feat_in += feat_in
        elif self.pool_mode == 'tap':
            self.feat_in = feat_in
        else:
            raise ValueError("pool mode for stats must be either tap or xvector based")

    def forward(self, encoder_output, length=None):
        mean = encoder_output.mean(dim=-1)  # Time Axis
        if self.pool_mode == 'xvector':
            std = encoder_output.std(dim=-1)
            pooled = torch.cat([mean, std], dim=-1)
        else:
            pooled = mean
        return pooled


seed_everything(42)

# Generate random encoder output
lengths = torch.tensor([40, 30, 20, 10])
batch_size = len(lengths)
feat_dim = 10
features_tensor = torch.zeros(size=[batch_size, feat_dim, lengths.max()])
features_list: List[torch.Tensor] = []
for i, length in enumerate(lengths):
    # Randomly shift mean and variance for each feature of each signal
    scale = torch.randn(size=[feat_dim, 1])
    bias = torch.randn(size=[feat_dim, 1])
    feats = torch.randn(size=[feat_dim, length]) / scale + bias
    features_list.append(feats)
    features_tensor[i, :, :length] = feats

old_pool: OldStatsPoolLayer = OldStatsPoolLayer(feat_in=feat_dim, pool_mode="xvector")
new_pool: StatsPoolLayer = StatsPoolLayer(feat_in=feat_dim, pool_mode="xvector")

old_stats: torch.Tensor = old_pool(encoder_output=features_tensor, length=lengths)
new_stats: torch.Tensor = new_pool(encoder_output=features_tensor, length=lengths)

print(f"Comparing batched statistics to individual stats")
for i in range(batch_size):
    # Extract stats with original stats pooler
    old_pooled_stats = old_stats[i]
    old_pool_mean = old_pooled_stats[:feat_dim]
    old_pool_std = old_pooled_stats[feat_dim:]
    # Extract stats with new stats pooler
    new_pooled_stats = new_stats[i]
    new_pool_mean = new_pooled_stats[:feat_dim]
    new_pool_std = new_pooled_stats[feat_dim:]
    # Compute the mean/std with PyTorch with individual signal (no padding)
    torch_mean = features_list[i].mean(dim=-1)
    torch_std = features_list[i].std(dim=-1)
    print(f"Batch element {i}, length {lengths[i].item()}:")
    print(f"\tOld mean all close: {torch.allclose(old_pool_mean, torch_mean)}")
    print(f"\tOld std all close: {torch.allclose(old_pool_std, torch_std)}")
    print(f"\tNew mean all close: {torch.allclose(new_pool_mean, torch_mean)}")
    print(f"\tNew std all close: {torch.allclose(new_pool_std, torch_std)}")
```

Expected output:

```text
Global seed set to 42
Comparing batched statistics to individual stats
Batch element 0, length 40:
        Old mean all close: True
        Old std all close: True
        New mean all close: True
        New std all close: True
Batch element 1, length 30:
        Old mean all close: False
        Old std all close: False
        New mean all close: True
        New std all close: True
Batch element 2, length 20:
        Old mean all close: False
        Old std all close: False
        New mean all close: True
        New std all close: True
Batch element 3, length 10:
        Old mean all close: False
        Old std all close: False
        New mean all close: True
        New std all close: True
```

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.