Training MirNetv2 and NAFNet in a Knowledge Distillation setup

[Samarendra109]

In the original paper (https://arxiv.org/abs/2204.04676) the authors find the optimal number of blocks to be 36. It has the minimal increase in latency and great increase in performance. So we need to train NAFNet with that setting. In the initial tests we see that NAFNet is able to beat MirnetV2 even with just 9 blocks (On LOL Dataset). So there is reason to believe a distilled NAFNet with 9 blocks trained from the teacher model with 36 blocks can give excellent performance. This needs to be tested.

[soumik12345]

TODOs:

• [ ] Build dataloaders for the following datasets:
  • [ ] SCIE dataset
  • [x] MIT Adobe 5K
  • [ ] See-In-Dark
  • [ ] Seeing Motion in the Dark
  • [ ] Seeing Dynamic Scene in the Dark
• [ ] Investigate if the aforementioned datasets can be combined to build one large dataset
• [ ] Train a large variant of MirNetv2 and NAFNet on this combined dataset.
• [ ] Distill into a smaller model:
  • [ ] Figure out a strategy to identify the layers the weights can be retained in the student model
  • [ ] Train the models in a teacher-student setup