SparK Implementation

[gabrielfruet]

Addresses #1462

Continues #1750

Inspired by: https://github.com/keyu-tian/SparK

https://arxiv.org/pdf/2301.03580

[IgorSusmelj]

FYI, you can update the branch with latest main. I should resolve the CI issues!

[liopeer]

@gabrielfruet I'll try to help a bit in the next few days so that we can merge this. Pushed a few cleanups and a context variable for taking care of that ugly global tensor. Feel free to also reach out on Discord for slightly faster interaction.

[gabrielfruet]

Honestly, i don't like the way SparK's author did this using globals.

An alternative would be to reuse the masks_per_level object that was already created and assign it to each Sparse converted layer.

We have two options here

1. Assign a new list container on each new forward pass with the new masks.
   • Would require to traverse the nn.Module each and every single time that the forward pass starts, identify sparse layers and assign the mask.
2. Have a "global" list that is assigned at construction time to the sparse layers with the len equals to the number of downsamples. (initially starting with null for example)
   • I think this wouldn't conflict with DDP and neither would need to traverse the module tree each time.
   • We should only make sure that the list is ONLY modified in-place, always. masks_per_level[level] = new_mask, such that the others can reuse it. The underlying memory address of the list would be the same for the lifetime of that module.

In any way both would work better as a design pattern.

Before i took any action here tell me WDYT. I think this could be a follow-up PR when this got merged. Or if you prefer, i can do this right away, just let me know your thoughts.

[liopeer]

I created a context variable now, which I think is a decent enough solution (minimal changes to the original code necessary, but a clear context instead of a global variable). I saw that you already picked it up and are using it, great! 👍
I can do another review probably tomorrow (Friday, 06.03.)

[gabrielfruet]

https://github.com/keyu-tian/SparK

This is the original repo. I joined some files into this single file and after that i started refactoring the classes and functions to adhere to the Lightly philosophy of providing the modules, not the full implementation.

[gabrielfruet]

Hi @liopeer, sorry for the delay. I got really sick recently.

I made some comments and reviews, but i forgot to submit them. I submitted them now.

[codecov]

Codecov Report

❌ Patch coverage is 33.55049% with 204 lines in your changes missing coverage. Please review.
✅ Project coverage is 84.00%. Comparing base (ee30cd4) to head (edbfdcf).
⚠️ Report is 4 commits behind head on master.

Files with missing lines	Patch %	Lines
lightly/models/modules/sparse_spark.py	27.17%	201 Missing ⚠️
lightly/loss/sparse_spark.py	89.65%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1892      +/-   ##
==========================================
- Coverage   86.10%   84.00%   -2.10%     
==========================================
  Files         168      171       +3     
  Lines        6979     7352     +373     
==========================================
+ Hits         6009     6176     +167     
- Misses        970     1176     +206     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[gabrielfruet]

@liopeer If you wish i could increase test coverage. If not necessary let me know.

[gabrielfruet]

I'll take a look into these type errors (Check 3.7)