Distillation Model PT
State: closed | Pull request: yes

# Changelog
- Prototype generic TeacherStudentModelPT

# Usage

Model extends a mixin class that enables registration of tensors inside of `training_step` for teacher student distillation loss as follows : 

```python
 # PTL-specific methods
    def training_step(self, batch, batch_nb):
        signal, signal_len, transcript, transcript_len = batch
        if isinstance(batch, DALIOutputs) and batch.has_processed_signal:
            log_probs, encoded_len, predictions = self.forward(
                processed_signal=signal, processed_signal_length=signal_len
            )
        else:
            log_probs, encoded_len, predictions = self.forward(input_signal=signal, input_signal_length=signal_len)

        loss_value = self.loss(
            log_probs=log_probs, targets=transcript, input_lengths=encoded_len, target_lengths=transcript_len
        )

        # Distillation support
        if self.is_being_distilled():
            if self.is_student_model():
                loss_key = 'input'
            else:
                loss_key = 'target'

            # Register the tensor for the loss function
            self.register_distillation_tensor(loss_key=loss_key, tensor=log_probs)
            # No need for further steps, return immediately since later elements are not available on
            # both the student and the teacher models
     
       ...
```

In any current NeMo model config, attach the following section

```yaml
model:
  # DISTILLATION DETAILS
  distillation:
    model_path: ???

    # loss weights
    distillation_loss_weight: 1.0
    student_train_loss_weight: 1.0
    cosine_loss_weight: 1.0

    # encoder matching
    distill_encoder: False  # attempts to distill encoder as well if model shapes are compatible

    # logits temperature
    temperature: 1.0

    # preserve memory for secondary loss computation (true by default)
    preserve_distillation_memory: false

    # optional - distillation loss config
    loss:
      primary:
        _target_: 'torch.nn.KLDivLoss'
        log_target: true
        reduction: 'batchmean'

      cosine:
        _target_: 'nemo.collections.common.losses.cosine_embedding.CosineEmbeddingLossWrapper'

  ... rest of student model config as usual
```

Then build a model using `DistillationModelPT` with this config and call `trainer.fit()`.

```python
import pytorch_lightning as pl
from omegaconf import OmegaConf

from nemo.core.classes.teacher_student import TeacherStudentModelPT
from nemo.core.config import hydra_runner
from nemo.utils import logging
from nemo.utils.exp_manager import exp_manager


@hydra_runner(config_path="configs/citrinet_bpe", config_name="citrinet_256_same_filters")
def main(cfg):
    # logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}')
    trainer = pl.Trainer(**cfg.trainer)
    exp_manager(trainer, cfg.get("exp_manager", None))

    teacher_student_model = DistillationModelPT(cfg=cfg.model, trainer=trainer)

    trainer.fit(teacher_student_model)


if __name__ == '__main__':
    main()

```

# Working
To be expanded once API is finalized.



This pull request **introduces 24 alerts** when merging fcbfbdd0a2934da73a8898d6688a615c4431cde4 into fbc5c2b4033896b75f643cf169976c1d38505d65 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-719061033791abc13da3e5efcb9ec3fd7f2232e4)

**new alerts:**

* 20 for Unused import
* 2 for Wrong name for an argument in a class instantiation
* 1 for First parameter of a method is not named &#39;self&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 24 alerts** when merging e92091de87d1e4075119d8ddbf6bbef9d4637f85 into fbc5c2b4033896b75f643cf169976c1d38505d65 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-6a7597d154e1828be7a6bbe16fb002d9cc3841c5)

**new alerts:**

* 20 for Unused import
* 2 for Wrong name for an argument in a class instantiation
* 1 for First parameter of a method is not named &#39;self&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 2 alerts** when merging 6f651e929b25df79861c72ab4b8d96218a600d49 into 56649822d824eba946f243e96dbe215996e03794 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-f7a24d2d58168e700c327d02e02381645de149a9)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 2 alerts** when merging 98e85727baeb2d04870dd9db09138a50ead45247 into 56649822d824eba946f243e96dbe215996e03794 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-6dfebbace6d0dc646a131844e5d0e94e672fc2a9)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 2 alerts** when merging 5f30bdccaf633e1b84f852c78b71ad90a66e6279 into e59dd39f6f847155a82d40ebe2264a0b7f2b7a13 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-0c7417a07a9c793e952a899be06612c3c6b86ae9)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 2 alerts** when merging 7205ac6bbaeeca0ff33dcada6a75ac576efdea36 into 627956e12c349fa9fde8a71d20d997ff5060294a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-faa1de9a1cb854f3c5d870c8d746db743ad2932a)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 2 alerts** when merging e21f4a98e065c6d7f153c6729c0f693695129630 into 627956e12c349fa9fde8a71d20d997ff5060294a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-bbb9751b5f5c51509c5275e10d6a24f5e8650943)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 2 alerts** when merging 52c07569485b56ae5aff1ad8c89d11c048ff8bc3 into 12a7e66dc90c5a604477ecb1ce37fee7320ba70b - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-91c475c72c7f16693ffcb3cd64837c1ad9cfe5a1)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 2 alerts** when merging d9ed96a77cffab46235b549192189b9952c3956e into 855ce265b80c0dc40f4f06ece76d2c9d6ca1be8d - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-1dc4d2efff9b42387291d6f1d2b769f3ae6b211b)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 3 alerts** when merging 1efebb387963cad9a49c5d1a3c35166d7ac4bf78 into 3d92dcf5ec6972b63c30faaff0aca4ffae16671d - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-146bd724d5ec4ab34fb2e82fa17017c65b15a294)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation
* 1 for Unused import

This pull request **introduces 3 alerts** when merging 681b73229178e16a05fb472a0df1dc2e74bfd3cd into 6d53452eb1487102c0d5ae657815ec2aecc12076 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-14d6d7ec52d383cd0f58db0a5327536297db43bb)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation
* 1 for Unused import

This pull request **introduces 5 alerts** when merging b5c9cc405503f2ae1c3ae5082af5f41df76f50b6 into f0b3624dc1aee9670d85b82251ad4e488a622bc7 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-0729e653242fddcd97df8ff03a9e7ca01192206b)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation
* 1 for Unused local variable

This pull request **introduces 5 alerts** when merging 26d8bbfaf04ba028f8f51a98e609c94397c119f0 into 08f3c65837adfcd092ad921c050ce1628cfcb8c4 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-84c4f8c0a8a7b64f84e05ea42872c57672675777)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation
* 1 for Unused local variable

This pull request **introduces 5 alerts** when merging 78900c51c5876b8edf3af889c60d77b0d4e98c29 into c036cec3acc9c2cfb39fbbead20deb6bbdde7b94 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-7fca2a7efd53a775693cc81dc790fd11884b2a75)

**new alerts:**

* 3 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 5 alerts** when merging 1da5af7556b585226c2c9f707adda831eb5c5d56 into c036cec3acc9c2cfb39fbbead20deb6bbdde7b94 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-6e931d7898ca2d3bd25df3a6700063f62e57fb19)

**new alerts:**

* 3 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 3 alerts** when merging b89addb530138e43e63eff16ff32da04eace035c into cfe2548d4db9dcc4ca393ee63604139c9f681cd8 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-5097ff851790a2e1b3bed6b7257d4e01517d71af)

**new alerts:**

* 2 for Wrong name for an argument in a class instantiation
* 1 for Unused import

This pull request **introduces 4 alerts** when merging b4d60a12ef604bad8a742961832c32a07fc1e399 into cfe2548d4db9dcc4ca393ee63604139c9f681cd8 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-f906ed361f365dc5590d965d4fb3e3b05576b74f)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 4 alerts** when merging 4b40a781139485255b49456d3c58594b5d89dbfd into e070e04ea219a53842bc6c3adb461a7b96c52519 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-04450883cae12e75669a7edc4f635fb5549c893a)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 4 alerts** when merging 27742c0940775f6a0e07710f3a3fa4a40cf0d309 into 5fcfa9e7bf2da39d622071fc0cac592f993f89ea - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-b10899532ace4825ee1546889c2643817917aa4f)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 4 alerts** when merging ce7f32161b2b7bfc157124b08902939f1b448e71 into 44a3d02b87692b45a1906972f63f127a6043c92b - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-7bf08e0a6bf7c379630a07f3bcb0f0c1aa612683)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 4 alerts** when merging 6c56055038d556beb14213f7774f1adac4917665 into a8b6a1a4dd0d89ea76e5c3ee5e75a8d9104700dd - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-20973e3f34bba46b58b04f049d727d8d3eb2937d)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation

Hi @titu1994 ,

Thanks a lot for this awesome work.
Do you have any idea when this pull request will be ready and merged to the main ?

It will take time. There are experiments underway to see the effectiveness of this approach. 

This pull request **introduces 4 alerts** when merging d14259af2a158beebd1f5fcd4c143f8005a0b956 into cd6a691607894f13cafcbe7994d1b2809ccb7aa4 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-316012550de5366a989699e74a0d486814898ba2)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 4 alerts** when merging 8cb9374293a596b86650a75a1d0e917da4716e75 into aad6bbf143331c43b08816ea316533fc6c98dc2a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-c037d3440a2ba8eed2860d4d689fcc5465201c71)

**new alerts:**

* 2 for Unused import
* 2 for Wrong name for an argument in a class instantiation

This pull request **introduces 1 alert** when merging 4fd7b9c09e9d41e48c855ee29838a5810d4090c3 into aad6bbf143331c43b08816ea316533fc6c98dc2a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-cb1f4e1db9bdfc3840bbad87c5a6f56210b63b06)

**new alerts:**

* 1 for Unused import

This pull request **fixes 1 alert** when merging 49fde6bcb30e119a852dfb46b7dbdf873e228f38 into aad6bbf143331c43b08816ea316533fc6c98dc2a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-0bd56d32d0eab057e6edd21947aa52dbaae6efd0)

**fixed alerts:**

* 1 for Unused import

This pull request **introduces 4 alerts** and **fixes 1** when merging 8225938eef1e1d01b0f8e8f65f18196fd8263127 into ec6591e76ab5d33e050e0dded25bed43d5d6c3e4 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-dff6fede60f9b2ce118407d88684cf8d43548f7f)

**new alerts:**

* 4 for Unused import

**fixed alerts:**

* 1 for Unused import

Closing due to lack of necessity at the moment.