Add files helpful for fine-tuning (#6)

- Update README.md (d59bc916fd05f54ebb9f907085cf619a3f11ab4b)
- Add some configs and a module (450a84c7cb080608946574878549592d92c24735)
- Add checkpoint compatible with GitHub repo's (6f594db60491d031b4c78786503ac9d8a7bc17e4)
- Add weights for DINO head (0f9f28c54f8b6d94e2666e050bffa676c897d59f)

README.md

@@ -58,6 +58,8 @@ Underlying biases of the training datasets may not be well characterized.
 
 ## Getting started
 
+### Get some data
+
 Let us first write an auxiliary function to download a chest X-ray.
 
 ```python
@@ -73,6 +75,8 @@ Let us first write an auxiliary function to download a chest X-ray.
 ...
 ```
 
+### Load the model
+
 Now let us download the model and encode an image.
 
 ```python
@@ -82,13 +86,17 @@ Now let us download the model and encode an image.
 >>>
 >>> # Download the model
 >>> repo = "microsoft/rad-dino"
->>> model = AutoModel.from_pretrained(repo)
+>>> rad_dino = AutoModel.from_pretrained(repo)
 >>>
 >>> # The processor takes a PIL image, performs resizing, center-cropping, and
 >>> # intensity normalization using stats from MIMIC-CXR, and returns a
 >>> # dictionary with a PyTorch tensor ready for the encoder
 >>> processor = AutoImageProcessor.from_pretrained(repo)
->>>
+```
+
+### Encode an image
+
+```python
 >>> # Download and preprocess a chest X-ray
 >>> image = download_sample_image()
 >>> image.size  # (width, height)
@@ -97,7 +105,7 @@ Now let us download the model and encode an image.
 >>>
 >>> # Encode the image!
 >>> with torch.inference_mode():
->>>     outputs = model(**inputs)
+>>>     outputs = rad_dino(**inputs)
 >>>
 >>> # Look at the CLS embeddings
 >>> cls_embeddings = outputs.pooler_output
@@ -124,6 +132,62 @@ We will use [`einops`](https://einops.rocks/) (install with `pip install einops`
 torch.Size([1, 768, 37, 37])
 ```
 
+### Weights for fine-tuning
+
+We have released a checkpoint compatible with
+[the original DINOv2 code](https://github.com/facebookresearch/dinov2) to help
+researchers fine-tune our model.
+
+First, let us write code to load a
+[`safetensors` checkpoint](https://huggingface.co/docs/safetensors).
+
+```python
+>>> import safetensors
+>>> def safetensors_to_state_dict(checkpoint_path: str) -> dict[str, torch.Tensor]:
+...     state_dict = {}
+...     with safe_open(checkpoint_path, framework="pt") as ckpt_file:
+...         for key in ckpt_file.keys():
+...             state_dict[key] = ckpt_file.get_tensor(key)
+...     return state_dict
+...
+```
+
+We can now use the hub model and load the RAD-DINO weights.
+Let's clone the DINOv2 repository so we can import the code for the head.
+
+```shell
+git clone https://github.com/facebookresearch/dinov2.git
+cd dinov2
+```
+
+```python
+>>> import torch
+>>> rad_dino_gh = torch.hub.load(".", "dinov2_vitb14")
+>>> backbone_state_dict = safetensors_to_state_dict("backbone_compatible.safetensors")
+>>> rad_dino_gh.load_state_dict(backbone_state_dict, strict=True)
+<All keys matched successfully>
+```
+
+The weights of the head are also released:
+
+```python
+>>> from dinov2.layers import DINOHead
+>>> rad_dino_head_gh = DINOHead(
+...    in_dim=768,
+...    out_dim=65536,
+...    hidden_dim=2048,
+...    bottleneck_dim=256,
+...    nlayers=3,
+... )
+>>> head_state_dict = safetensors_to_state_dict("dino_head.safetensors")
+>>> rad_dino_head_gh.load_state_dict(head_state_dict, strict=True)
+<All keys matched successfully>
+```
+
+### Configs and augmentation
+
+The configuration files [`ssl_default_config.yaml`](./ssl_default_config.yaml) and [`vitb14_cxr.yaml`](./vitb14_cxr.yaml), and the [`augmentations`](./augmentations.py) module are also available in the repository to help researchers reproduce the training procedure with our hyperparameters.
+
 ## Training details
 
 ### Training data
@@ -225,4 +289,4 @@ We used [SimpleITK](https://simpleitk.org/) and [Pydicom](https://pydicom.github
 
 ## Model card contact
 
-Fernando Pérez-García ([`[email protected]`](mailto:[email protected])).
+Fernando Pérez-García ([`[email protected]`](mailto:[email protected])).

augmentations.py

@@ -0,0 +1,147 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# See LICENSE in the repo root for license information.
+#
+# Portions:
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import logging
+
+from PIL import Image
+from torchvision import transforms
+
+from .transforms import (
+    GaussianBlur,
+    MaybeToTensor,
+    make_normalize_transform,
+)
+
+
+logger = logging.getLogger("dinov2")
+
+
+class DataAugmentationDINO(object):
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        global_crops_size=224,
+        local_crops_size=96,
+    ):
+        self.global_crops_scale = global_crops_scale
+        self.local_crops_scale = local_crops_scale
+        self.local_crops_number = local_crops_number
+        self.global_crops_size = global_crops_size
+        self.local_crops_size = local_crops_size
+
+        logger.info("###################################")
+        logger.info("Using data augmentation parameters:")
+        logger.info(f"global_crops_scale: {global_crops_scale}")
+        logger.info(f"local_crops_scale: {local_crops_scale}")
+        logger.info(f"local_crops_number: {local_crops_number}")
+        logger.info(f"global_crops_size: {global_crops_size}")
+        logger.info(f"local_crops_size: {local_crops_size}")
+        logger.info("###################################")
+
+        # random resized crop and flip
+        self.geometric_augmentation_global = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    global_crops_size, scale=global_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC
+                ),
+                transforms.RandomHorizontalFlip(p=0.5),
+            ]
+        )
+
+        self.geometric_augmentation_local = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    local_crops_size, scale=local_crops_scale, interpolation=transforms.InterpolationMode.BICUBIC
+                ),
+                transforms.RandomHorizontalFlip(p=0.5),
+            ]
+        )
+
+        # color distorsions / blurring
+        color_jittering = transforms.Compose(
+            [
+                transforms.RandomApply(
+                    [transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)],
+                    p=0.8,
+                ),
+                transforms.RandomGrayscale(p=0.2),
+            ]
+        )
+
+        global_transfo1_extra = GaussianBlur(p=0.5)
+
+        global_transfo2_extra = transforms.Compose(
+            [
+                GaussianBlur(p=0.1),
+            ]
+        )
+
+        local_transfo_extra = GaussianBlur(p=0.5)
+
+        # normalization
+        self.normalize = transforms.Compose(
+            [
+                MaybeToTensor(),
+                make_normalize_transform(),
+            ]
+        )
+
+        self.global_transfo1 = transforms.Compose([color_jittering, global_transfo1_extra, self.normalize])
+        self.global_transfo2 = transforms.Compose([color_jittering, global_transfo2_extra, self.normalize])
+        self.local_transfo = transforms.Compose([color_jittering, local_transfo_extra, self.normalize])
+
+    def __call__(self, image):
+        output = {}
+
+        # global crops:
+        im1_base = self.geometric_augmentation_global(image)
+        global_crop_1 = self.global_transfo1(im1_base)
+
+        im2_base = self.geometric_augmentation_global(image)
+        global_crop_2 = self.global_transfo2(im2_base)
+
+        output["global_crops"] = [global_crop_1, global_crop_2]
+
+        # global crops for teacher:
+        output["global_crops_teacher"] = [global_crop_1, global_crop_2]
+
+        # local crops:
+        local_crops = [
+            self.local_transfo(self.geometric_augmentation_local(image)) for _ in range(self.local_crops_number)
+        ]
+        output["local_crops"] = local_crops
+        output["offsets"] = ()
+
+        return output
+
+
+def get_online_classification_augmentation_from_config(cfg) -> transforms.Compose:
+    augmentation_config = cfg.evaluation.online.augmentation
+    interpolation = getattr(Image.Resampling, augmentation_config.interpolation)
+    resize_size = crop_size = cfg.crops.global_crops_size
+    resize = transforms.Resize(resize_size, interpolation=interpolation)
+    crop = transforms.CenterCrop(crop_size)
+    affine = transforms.RandomAffine(
+        degrees=augmentation_config.degrees,
+        scale=augmentation_config.scale,
+        shear=augmentation_config.shear,
+        interpolation=interpolation,
+    )
+    transforms_list = [
+        resize,
+        crop,
+        affine,
+        MaybeToTensor(),
+        make_normalize_transform(),
+    ]
+    if augmentation_config.horizontal_flip:
+        transforms_list.append(transforms.RandomHorizontalFlip())
+    return transforms.Compose(transforms_list)

backbone_compatible.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1eac0464b2a00d368aa3eea1dc029964b10320fbabc59a8a4e768c43a23d26f4
+size 346338024

dino_head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b3599663464ed1054f7777f547db02f518581acc5becdd3eddffc8c507f3778
+size 92554920

ssl_default_config.yaml

@@ -0,0 +1,135 @@
+MODEL:
+  WEIGHTS: ''
+compute_precision:
+  grad_scaler: true
+  teacher:
+    backbone:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+    dino_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+    ibot_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+  student:
+    backbone:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp16
+        buffer_dtype: fp32
+    dino_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp32
+        buffer_dtype: fp32
+    ibot_head:
+      sharding_strategy: SHARD_GRAD_OP
+      mixed_precision:
+        param_dtype: fp16
+        reduce_dtype: fp32
+        buffer_dtype: fp32
+dino:
+  loss_weight: 1.0
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_nlayers: 3
+  head_hidden_dim: 2048
+  koleo_loss_weight: 0.1
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  separate_head: false
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_nlayers: 3
+  head_hidden_dim: 2048
+train:
+  batch_size_per_gpu: 64
+  dataset_path: ImageNet:split=TRAIN
+  output_dir: .
+  saveckp_every_n_epoch: 5
+  seed: 0
+  num_workers: 10
+  OFFICIAL_EPOCH_LENGTH: 0 # automatic rescaling based on the dataset len is applied if this is set to 0
+  cache_dataset: true
+  centering: "centering" # or "sinkhorn_knopp"
+student:
+  arch: vit_large
+  patch_size: 16
+  drop_block_rate: 0.0
+  drop_path_rate: 0.3
+  layerscale: 1.0e-05
+  drop_path_uniform: true
+  pretrained_weights: ''
+  ffn_layer: "mlp"
+  block_chunks: 0
+  qkv_bias: true
+  proj_bias: true
+  ffn_bias: true
+  num_register_tokens: 0
+  interpolate_antialias: false
+  interpolate_offset: 0.1
+  load_weights: true
+  checkpoints_dir: null
+teacher:
+  momentum_teacher: 0.992
+  final_momentum_teacher: 1
+  warmup_teacher_temp: 0.04
+  teacher_temp: 0.07
+  warmup_teacher_temp_epochs: 30
+optim:
+  epochs: 100
+  weight_decay: 0.04
+  weight_decay_end: 0.4
+  base_lr: 0.004  # learning rate for a batch size of 1024
+  lr: 0.  # will be set after applying scaling rule
+  warmup_epochs: 10
+  min_lr: 1.0e-06
+  clip_grad: 3.0
+  freeze_last_layer_epochs: 1
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  layerwise_decay: 0.9
+  adamw_beta1: 0.9
+  adamw_beta2: 0.999
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 224
+  local_crops_size: 96
+evaluation:
+  eval_period_iterations: 12500
+  dataset_str: None
+  online:  # see dinov2.eval.linear_callback for documentation
+    learning_rate: 1e-6  # will be multiplied by batch size and number of devices
+    num_last_blocks: 1
+    add_avg_pool: true
+    num_update_epochs_per_eval: 3
+    augmentation:
+      degrees: 30
+      scale:
+      - 0.8
+      - 1.2
+      shear: 15
+      interpolation: BICUBIC
+      horizontal_flip: true

vitb14_cxr.yaml

@@ -0,0 +1,31 @@
+# this corresponds to the CXR config
+train:
+  batch_size_per_gpu: 40  # For nodes with v100s (32 GB), use 20.
+  saveckp_every_n_epoch: 25
+student:
+  arch: vit_base
+  block_chunks: 4
+  patch_size: 14
+  drop_block_rate: 0.00
+  drop_path_rate: 0.30
+teacher:
+  warmup_teacher_temp_epochs: 50
+optim:
+  epochs: 100
+  warmup_epochs: 5
+  base_lr: 0.001
+evaluation:
+  eval_period_iterations: 300
+  tasks:  # from the metadata.csv file of the CANDID processed dataset
+  - pneumothorax
+crops:
+    global_crops_size: 518
+    local_crops_size: 196
+    global_crops_scale:
+    - 0.50
+    - 1.00
+    local_crops_number: 8
+    local_crops_scale:
+    - 0.20
+    - 0.50
+pretrained: true