.. currentmodule:: slideflow.simclr

.. _simclr_ssl:

Self-Supervised Learning (SSL)

==============================

Slideflow provides easy access to training the self-supervised, contrastive learning framework `SimCLR <https://arxiv.org/abs/2002.05709>`_. Self-supervised learning provides an avenue for learning useful visual representations in your dataset without requiring ground-truth labels. These visual representations can be exported as feature vectors and used for downstream analyses such as :ref:`dimensionality reduction <slidemap>` or :ref:`multi-instance learning <mil>`.

The ``slideflow.simclr`` module contains a `forked Tensorflow implementation <https://github.com/jamesdolezal/simclr/>`_ minimally modified to interface with Slideflow. SimCLR models can be trained with :meth:`slideflow.Project.train_simclr`, and SimCLR features can be calculated as with other models using :meth:`slideflow.Project.generate_features`.

Training SimCLR

***************

First, determine the SimCLR training parameters with :func:`slideflow.simclr.get_args`. This function accepts parameters via keyword arguments, such as ``learning_rate`` and ``temperature``, and returns a configured :class:`slideflow.simclr.SimCLR_Args`.

.. code-block:: python

    from slideflow import simclr

    args = simclr.get_args(

        temperature=0.1,

        learning_rate=0.3,

        train_epochs=100,

        image_size=299

    )

Next, assemble a training and (optionally) a validation dataset. The validation dataset is used to assess contrastive loss during training, but is not required.

.. code-block:: python

    import slideflow as sf

    # Load a project and dataset

    P = sf.load_project('path')

    dataset = P.dataset(tile_px=299, tile_um=302)

    # Split dataset into training/validation

    train_dts, val_dts = dataset.split(

        val_fraction=0.3,

        model_type='classification',

        labels='subtype')

Finally, SimCLR can be trained with :meth:`slideflow.Project.train_simclr`. You can train with a single dataset:

.. code-block:: python

    P.train_simclr(args, dataset)

You can train with an optional validation dataset:

.. code-block:: python

    P.train_simclr(

        args,

        train_dataset=train_dts,

        val_dataset=val_dts

    )

And you can also optionally provide labels for training the supervised head. To train a supervised head, you'll also need to set the SimCLR argument ``lineareval_while_pretraining=True``.

.. code-block:: python

    # SimCLR args

    args = simclr.get_args(

        ...,

        lineareval_while_pretraining=True

    )

    # Train with validation & supervised head

    P.train_simclr(

        args,

        train_dataset=train_dts,

        val_dataset=val_dts,

        outcomes='subtype'

    )

The SimCLR model checkpoints and final saved model will be saved in the ``simclr/`` folder within the project root directory.

.. _dinov2:

Training DINOv2

***************

A lightly modified version of `DINOv2 <https://arxiv.org/abs/2304.07193>`__ with Slideflow integration is available on `GitHub <https://github.com/jamesdolezal/dinov2>`_. This version facilitates training DINOv2 with Slideflow datasets and adds stain augmentation to the training pipeline.

To train DINOv2, first install the package:

.. code-block:: bash

    pip install git+https://github.com/jamesdolezal/dinov2.git

Next, configure the training parameters and datsets by providing a configuration YAML file. This configuration file should contain a ``slideflow`` section, which specifies the Slideflow project and dataset to use for training. An example YAML file is shown below:

.. code-block:: yaml

    train:

      dataset_path: slideflow

      batch_size_per_gpu: 32

      slideflow:

        project: "/mnt/data/projects/TCGA_THCA_BRAF"

        dataset:

          tile_px: 299

          tile_um: 302

          filters:

            brs_class:

            - "Braf-like"

            - "Ras-like"

        seed: 42

        outcome_labels: "brs_class"

        normalizer: "reinhard_mask"

        interleave_kwargs: null

See the `DINOv2 README <https://github.com/jamesdolezal/dinov2>`_ for more details on the configuration file format.

Finally, train DINOv2 using the same command-line interface as the original DINOv2 implementation. For example, to train DINOv2 on 4 GPUs on a single node:

.. code-block:: bash

    torchrun --nproc_per_node=4 -m "dinov2.train.train" \

        --config-file /path/to/config.yaml \

        --output-dir /path/to/output_dir

The teacher weights will be saved in ``outdir/eval/.../teacher_checkpoint.pth``, and the final configuration YAML will be saved in ``outdir/config.yaml``.

Generating features

*******************

Generating features from a trained SSL is straightforward - use the same :meth:`slideflow.Project.generate_features` and :class:`slideflow.DatasetFeatures` interfaces as :ref:`previously described <dataset_features>`, providing a path to a saved SimCLR model or checkpoint.

.. code-block:: python

    import slideflow as sf

    # Create the SimCLR feature extractor

    simclr = sf.build_feature_extractor(

        'simclr',

        ckpt='/path/to/simclr.ckpt'

    )

    # Calculate SimCLR features for a dataset

    features = P.generate_features(simclr, ...)

For DINOv2 models, use ``'dinov2'`` as the first argument, and pass the model configuration YAML file to ``cfg`` and the teacher checkpoint weights to ``weights``.

.. code-block:: python

    dinov2 = build_feature_extractor(

        'dinov2',

        weights='/path/to/teacher_checkpoint.pth',

        cfg='/path/to/config.yaml'

    )