Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Single-cell phenotyping with contrastive learning #113

Merged
merged 63 commits into from
Aug 2, 2024

Conversation

ziw-liu
Copy link
Collaborator

@ziw-liu ziw-liu commented Jul 19, 2024

Self-supervised learning of single-cell image patches with contrastive triplet loss.

mattersoflight and others added 30 commits June 30, 2024 22:06
* initial dataloader.py

* Update dataloader_test.py

* Update dataloader_test.py

* Update dataloader_test.py

* Update dataloader_test.py

* rename training script
@mattersoflight
Copy link
Member

@ziw-liu looks like you are powering through the refactor. Awesome.

Couple of notes:

  • Please integrate the current contrastive learning training script into the CLI, via a required flag like --model OR via a config field for the lighting module. The latter may be a seamless and low-maintenance route as we implement lightning modules for tasks like semantic segmentation.
  • applications/contrastive_phenotyping/dataloader_test.py can be moved to viscy/scripts. It can also be improved in the future PR to write a CLI for testing and logging parallelized data loading for given configuration.

@ziw-liu ziw-liu linked an issue Jul 23, 2024 that may be closed by this pull request
8 tasks
@ziw-liu ziw-liu added the enhancement New feature or request label Jul 29, 2024
ziw-liu and others added 10 commits July 29, 2024 11:43
* updated prediction code

* updated predict code

* updated code

* fixed the stem and forward pass (#115)

* fixed the stem and forward pass

* update forward calls to encoder

* self.encoder -> self.model

* nits

* l2 normalize projections

* black compliance

* black compliance

* WIP: Save progress before merging

* updated contrastive.py

* stem update

* updated predict code

* Delete viscy/applications/contrastive_phenotyping/PCA.ipynb

* pushing dataloader test updated

* pca deleted

* training and dataloader test

* updated structure

* deleted files

* updated training merged files

* removed commented code

* removed uneeded code

* removed uneeded code

* removed comments

* snake_case

* fixed CI issues

* removed num_fovs

---------

Co-authored-by: Shalin Mehta <[email protected]>
@ziw-liu ziw-liu marked this pull request as ready for review July 31, 2024 23:15
@ziw-liu
Copy link
Collaborator Author

ziw-liu commented Jul 31, 2024

@mattersoflight does #116 need to be fixed in this PR?

Copy link
Member

@mattersoflight mattersoflight left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ziw-liu #116 can be handled in the next PR. I will be able to review that PR. If @alishbaimran is able to train models with this branch, please merge.

Please also clean up multiple versions of code for constructing stem.

@@ -65,7 +65,7 @@ def _get_convnext_stage(


class UNeXt2Stem(nn.Module):
"""Stem for UNeXt2 networks."""
"""Stem for UNeXt2 and ContrastiveEncoder networks."""
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ziw-liu I wrote the StemDepthtoChannels to work for resnet and convnext models. Can you test the construction of UNeXt2 model with this stem and report back if the construction works? If it does, we should deprecate other classes for building the stem to avoid confusion in the future.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should delete this and have only one (flexible) implementation of stem.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ziw-liu ziw-liu mentioned this pull request Aug 2, 2024
@ziw-liu ziw-liu merged commit 0a72176 into main Aug 2, 2024
4 checks passed
@ziw-liu ziw-liu deleted the contrastive_phenotyping branch August 2, 2024 17:16
edyoshikun pushed a commit that referenced this pull request Aug 7, 2024
* first draft of contrastive learning model

* fixed stem and projection head, drafted lightning module

* Contrastive_dataloader (#99)

* initial dataloader.py

* Update dataloader_test.py

* Update dataloader_test.py

* Update dataloader_test.py

* Update dataloader_test.py

* rename training script

* move contrastive network to viscy.representation module

* Update hcs.py

* refactored class names

* correct imports

* cleaner names for model arch and module

* new imports

* Fixed epoch loss logging and WandB integration in ContrastiveModule

* updated training_script.py

* Update hcs.py

* contrastive.py

* engine.py

* script to test data i/o speed from different filesystems

* moved applications folder to viscy.applications so that pip install -e . works.

* add resnet50 to ContrastiveEncoder

* rename training_script.py to training_script_resnet.py

* test dataloader on lustre and vast

* move training_script_resnet to viscy.applications so that `pip install -e .` works

* refined the tests for contrastive dataloader

* sbatch script for dataloader

* delete redundant module

* nits: updated the model construction of contrastive resnet encoder.

* Updated training script, HCS data handling, engine, and contrastive representation

* Fix normalization, visualization issues, logging and multi-channel prediction

* updated training and prediction

* update training and prediction script

* formatting

* combine the application directories

* lint

* replace notebook with script

* format script

* rename scripts conflicting with pytest

* lint application scripts

* do not filter all warnings

* log instead of print

* split data modules by task

* clean up imports

* update typing

* use pathlib

* remove redundant file

* updated predict.py

* better typing

* wip: triplet dataset

* avoid forward ref
this might increase code analysis time a tiny bit
but should not have any effect at runtime

* check that z range is valid
and fix indexing

* clean up and explain random sampling

* sample dict instead of tuple and include track index

* take out generic HCS methods for reuse

* implement TripletDataModule

* use new batch type in engine

* better typing

* read normalization metadata

* docstring for data module

* drop normalization metadata after transformation

* remove unused import

* fix initial crop size

* Infection state (#118)

* updated prediction code

* updated predict code

* updated code

* fixed the stem and forward pass (#115)

* fixed the stem and forward pass

* update forward calls to encoder

* self.encoder -> self.model

* nits

* l2 normalize projections

* black compliance

* black compliance

* WIP: Save progress before merging

* updated contrastive.py

* stem update

* updated predict code

* Delete viscy/applications/contrastive_phenotyping/PCA.ipynb

* pushing dataloader test updated

* pca deleted

* training and dataloader test

* updated structure

* deleted files

* updated training merged files

* removed commented code

* removed uneeded code

* removed uneeded code

* removed comments

* snake_case

* fixed CI issues

* removed num_fovs

---------

Co-authored-by: Shalin Mehta <[email protected]>

---------

Co-authored-by: Shalin Mehta <[email protected]>
Co-authored-by: Alishba Imran <[email protected]>
Co-authored-by: Alishba Imran <[email protected]>
Co-authored-by: Alishba Imran <[email protected]>
Co-authored-by: Duo Peng <[email protected]>
edyoshikun added a commit that referenced this pull request Aug 8, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

WIP: refactor contrastive learning code with virtual staining code
4 participants