Upload 13 files

Dockerfile

@@ -0,0 +1,62 @@
+FROM nvidia/cuda:11.0-base-ubuntu20.04
+
+ENV DEBIAN_FRONTEND noninteractive
+
+ENV CUDNN_VERSION=8.0.5.39-1+cuda11.1
+ENV NCCL_VERSION=2.7.8-1+cuda11.1
+
+ARG python=3.8
+ENV PYTHON_VERSION=${python}
+
+# Set default shell to /bin/bash
+SHELL ["/bin/bash", "-cu"]
+
+RUN apt-get update && apt-get install -y --allow-downgrades \
+    --allow-change-held-packages --no-install-recommends \
+    build-essential \
+    cmake \
+    git \
+    curl \
+    vim \
+    wget \
+    ca-certificates \
+    libcudnn8=${CUDNN_VERSION} \
+    libnccl2=${NCCL_VERSION} \
+    libnccl-dev=${NCCL_VERSION} \
+    libjpeg-dev \
+    libpng-dev \
+    python${PYTHON_VERSION} \
+    python${PYTHON_VERSION}-dev \
+    python${PYTHON_VERSION}-distutils \
+    librdmacm1 \
+    libibverbs1 \
+    ibverbs-providers
+
+RUN ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
+
+RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
+    python get-pip.py && \
+    rm get-pip.py
+
+RUN /usr/bin/python -m pip install --upgrade pip
+
+# Install pytorch
+RUN pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 \
+-f https://download.pytorch.org/whl/torch_stable.html 
+
+RUN pip install tensorboard==2.5.0
+RUN pip install tensorboard-data-server==0.6.1
+RUN pip install tensorboard-plugin-wit==1.8.0
+RUN pip install tensorboardX==1.8
+
+RUN pip install timm==0.4.5
+RUN pip install opencv-contrib-python-headless==4.5.2.54
+RUN pip install tqdm==4.61.2
+RUN pip install PyYAML==5.4.1
+RUN pip install Pillow==8.3.1
+RUN pip install einops==0.3.0
+RUN pip install scipy==1.7.1
+
+
+
+

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2022 Garena Online Private Limited
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,8 +1,180 @@
----
-license: apache-2.0
-datasets:
-- imagenet-1k
-metrics:
-- accuracy
-pipeline_tag: image-classification
----
+# Mugs: A Multi-Granular Self-Supervised Learning Framework
+
+This is a PyTorch implementation of **Mugs** proposed by our paper "**Mugs: A Multi-Granular Self-Supervised Learning Framework**". [![arXiv](https://img.shields.io/badge/arXiv-2203.14415-b31b1b.svg?style=flat)](http://arxiv.org/abs/2203.14415)
+	
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/mugs-a-multi-granular-self-supervised/self-supervised-image-classification-on)](https://paperswithcode.com/sota/self-supervised-image-classification-on?p=mugs-a-multi-granular-self-supervised)
+
+<div align="center">
+<img width="100%" alt="Overall framework of Mugs. " src="./exp_illustration/framework.png">
+</div>
+
+**<p align="center">Fig 1. Overall framework of Mugs.** In (a), for each image, two random crops of one image 
+are fed into backbones of student and teacher. Three granular supervisions: 1) instance discrimination supervision, 2) local-group discrimination 
+supervision, and 3) group discrimination supervision, are adopted to learn multi-granular representation. In (b), local-group modules in 
+student/teacher averages all patch tokens, and finds top-k neighbors from memory buffer to aggregate them with the average for obtaining a local-group feature.</p>
+
+
+
+
+
+# Pretrained models on ImageNet-1K
+
+You can choose to download only the weights of the pretrained backbone used for downstream tasks, or the full checkpoint which contains backbone and projection head weights for both student and teacher networks.
+**<p align="center">Table 1. KNN and linear probing performance with their corresponding hyper-parameters, logs and model weights.</p>**
+<table>
+  <tr>
+    <th>arch</th>
+    <th>params</th>
+    <th>pretraining epochs</th>
+    <th>k-nn</th>
+    <th>linear</th>
+    <th colspan="6">download</th>
+  </tr>
+  <tr>
+    <td>ViT-S/16</td>
+    <td>21M</td>
+    <td>100</td>
+    <td>72.3%</td>
+    <td>76.4%</td>
+    <td><a href="https://drive.google.com/file/d/1V2TyArzr7qY93UFglPBHRfYVyAMEfsHR/view?usp=sharing">backbone only</a></td>
+    <td><a href="https://drive.google.com/file/d/1AePcCeUEhK0nb9syQKufqqnhpUEr9Rji/view?usp=sharing">full ckpt</a></td>
+    <td><a href="https://drive.google.com/file/d/17phHQx88f4_xSqkPtIYvUoUE2U-1tovg/view?usp=sharing">args</a></td>
+    <td><a href="https://drive.google.com/file/d/1UBMTB-C3BnNKT5939fhSstHc9H30Vizd/view?usp=sharing">logs</a></td>
+    <td><a href="https://drive.google.com/file/d/1MkXctkgqEXjWWRs4Cz5CyTTx_IHDOP4G/view?usp=sharing">eval logs</a></td>
+  </tr>
+  <tr>
+    <td>ViT-S/16</td>
+    <td>21M</td>
+    <td>300</td>
+    <td>74.8%</td>
+    <td>78.2%</td>
+    <td><a href="https://drive.google.com/file/d/1ZAPQ0HiDZO5Uk7jVqF46H6VbGxunZkuf/view?usp=sharing">backbone only</a></td>
+    <td><a href="https://drive.google.com/file/d/1EO-_kYlAt23qgFYZF2u-KLks5js9LvrZ/view?usp=sharing">full ckpt</a></td>
+    <td><a href="https://drive.google.com/file/d/1b6zLZ3r_mZbk17SvhJIZF2VCoYVbJUnU/view?usp=sharing">args</a></td>
+    <td><a href="https://drive.google.com/file/d/1L7VzH1rztoraBCBNVWL-Y8k7Y8PFU773/view?usp=sharing">logs</a></td>
+    <td><a href="https://drive.google.com/file/d/1KgnX8ReXIVsu65_-p7NWPH8S0HEDPMUU/view?usp=sharing">eval logs</a></td>
+  </tr>
+  <tr>
+    <td>ViT-S/16</td>
+    <td>21M</td>
+    <td>800</td>
+    <td>75.6%</td>
+    <td>78.9%</td>
+    <td><a href="https://drive.google.com/file/d/1KMdhxxWc2JXAiFqVxX584V4RvlJgckGq/view?usp=sharing">backbone only</a></td>
+    <td><a href="https://drive.google.com/file/d/1FBaOt0Rjxm6yyJadttOyN6hSh8ueZ0dh/view?usp=sharing">full ckpt</a></td>
+    <td><a href="https://drive.google.com/file/d/19Ma-eSIgdwLoBg6wBXeFiW46zCI2EHvH/view?usp=sharing">args</a></td>
+    <td><a href="https://drive.google.com/file/d/1wX4AUO5NBVZUb8jN1iGBRkS17sszb4_O/view?usp=sharing">logs</a></td>
+    <td><a href="https://drive.google.com/file/d/12tiO4glWZNB044TYiPPCfbnUX_9AbqVc/view?usp=sharing">eval logs</a></td>
+  </tr>
+  <tr>
+    <td>ViT-B/16</td>
+    <td>85M</td>
+    <td>400</td>
+    <td>78.0%</td>
+    <td>80.6%</td>
+    <td><a href="https://drive.google.com/file/d/13NUziwToBXBmS7n7V_1Z5N6EG_7bcncW/view?usp=sharing">backbone only</a></td>
+    <td><a href="https://drive.google.com/file/d/1M41TVVFyVRDTK5kbgLCEImrxw0AVtebb/view?usp=sharing">full ckpt</a></td>
+    <td><a href="https://drive.google.com/file/d/1-5fB5ZCVQAfxTXZ6ro56AVkhb3whpaJc/view?usp=sharing">args</a></td>
+    <td><a href="https://drive.google.com/file/d/11RlCx6eViRnFD6gBlr_lOOxOhu-L6l6D/view?usp=sharing">logs</a></td>
+    <td><a href="https://drive.google.com/file/d/1gOR250QFLZfe40pLNPcOqaLPAnKLuE_C/view?usp=sharing">eval logs</a></td>
+  </tr>
+  <tr>
+    <td>ViT-L/16</td>
+    <td>307M</td>
+    <td>250</td>
+    <td>80.3%</td>
+    <td>82.1%</td>
+    <td><a href="https://drive.google.com/file/d/1K76a-YnFYcmDXUZ_UlYVYFrWOt2a6733/view?usp=sharing">backbone only</a></td>
+    <td><a href="https://drive.google.com/file/d/1Q5Ukvucx44YawyOhMEAY13Ppb8OOWOAB/view?usp=sharing">full ckpt</a></td>
+    <td><a href="https://drive.google.com/file/d/1p8XhaA2_Zbejm__UT8iNKG8r5tzS9c6c/view?usp=sharing">args</a></td>
+    <td><a href="https://drive.google.com/file/d/1JLVcUNfkyBI0BcMm7OpNU_3KTxIABK0Z/view?usp=sharing">logs</a></td>
+    <td><a href="https://drive.google.com/file/d/1rqWenRFN0czat_55GY9GNOu7gS6fww3g/view?usp=sharing">eval logs</a></td>
+  </tr>
+</table>
+
+<div align="center">
+  <img width="100%" alt="Comparison of linear probing accuracy on ImageNet-1K." src="./exp_illustration/comparison.png">
+</div>
+
+**<p align="center">Fig 2. Comparison of linear probing accuracy on ImageNet-1K.**</p>
+
+## Pretraining Settings
+
+### Environment
+For reproducing, please install [PyTorch](https://pytorch.org/) and download the [ImageNet](https://imagenet.stanford.edu/) dataset.
+This codebase has been developed with python version 3.8, PyTorch version 1.7.1, CUDA 11.0 and torchvision 0.8.2. For the full 
+environment, please refer to our `Dockerfile` file. 
+
+
+### ViT pretraining :beer:
+To pretraining each model, please find the exact hyper-parameter settings at the `args` column of  [Table 1](https://github.com/sail-sg/mugs). For training log and linear probing log, please refer to the 
+`log` and `eval logs` column of [Table 1](https://github.com/sail-sg/mugs). 
+
+#### ViT-Small pretraining:
+To run ViT-small for 100 epochs, we use two nodes of total 8 A100 GPUs (total 512 minibatch size) by using following command: 
+```
+python -m torch.distributed.launch --nproc_per_node=8 main.py --data_path DATASET_ROOT --output_dir OUTPUT_ROOT --arch vit_small 
+--group_teacher_temp 0.04 --group_warmup_teacher_temp_epochs 0 --weight_decay_end 0.2 --norm_last_layer false --epochs 100
+```
+To run ViT-small for 300 epochs, we use two nodes of total 16 A100 GPUs (total 1024 minibatch size) by using following command:
+```
+python -m torch.distributed.launch --nproc_per_node=16 main.py --data_path DATASET_ROOT --output_dir OUTPUT_ROOT --arch vit_small 
+--group_teacher_temp 0.07 --group_warmup_teacher_temp_epochs 30 --weight_decay_end 0.1 --norm_last_layer false --epochs 300
+```
+To run ViT-small for 800 epochs, we use two nodes of total 16 A100 GPUs (total 1024 minibatch size) by using following command:
+```
+python -m torch.distributed.launch --nproc_per_node=16 main.py --data_path DATASET_ROOT --output_dir OUTPUT_ROOT --arch vit_small 
+--group_teacher_temp 0.07 --group_warmup_teacher_temp_epochs 30 --weight_decay_end 0.1 --norm_last_layer false --epochs 800
+```
+
+#### ViT-Base pretraining:
+To run ViT-base for 400 epochs, we use two nodes of total 24 A100 GPUs (total 1024 minibatch size) by using following command:
+```
+python -m torch.distributed.launch --nproc_per_node=24 main.py --data_path DATASET_ROOT --output_dir OUTPUT_ROOT --arch vit_base 
+--group_teacher_temp 0.07 --group_warmup_teacher_temp_epochs 50 --min_lr 2e-06 --weight_decay_end 0.1 --freeze_last_layer 3 --norm_last_layer 
+false --epochs 400
+```
+
+#### ViT-Large pretraining:
+To run ViT-large for 250 epochs, we use two nodes of total 40 A100 GPUs (total 640 minibatch size) by using following command:
+```
+python -m torch.distributed.launch --nproc_per_node=40 main.py --data_path DATASET_ROOT --output_dir OUTPUT_ROOT --arch vit_large 
+--lr 0.0015 --min_lr 1.5e-4 --group_teacher_temp 0.07 --group_warmup_teacher_temp_epochs 50 --weight_decay 0.025 
+--weight_decay_end 0.08 --norm_last_layer true --drop_path_rate 0.3 --freeze_last_layer 3 --epochs 250
+```
+
+## Evaluation
+We are cleaning up the evalutation code and will release them when they are ready.
+
+## Self-attention visualization
+Here we provide the self-attention map of the [CLS] token on the heads of the last layer 
+<div align="center">
+  <img width="100%" alt="Self-attention from a ViT-Base/16 trained with Mugs" src="./exp_illustration/attention_vis.png">
+</div>
+
+**<p align="center">Fig 3. Self-attention from a ViT-Base/16 trained with Mugs.**</p>
+
+
+## T-SNE visualization
+Here we provide the T-SNE visualization of the learned feature by ViT-B/16. 
+We show the fish classes in ImageNet-1K, i.e., the first six classes, 
+including tench, goldfish, white shark, tiger shark, hammerhead, electric 
+ray. See more examples in Appendix.
+<div align="center">
+  <img width="100%" alt="T-SNE visualization of the learned feature by ViT-B/16." src="./exp_illustration/TSNE.png">
+</div>
+
+**<p align="center">Fig 4. T-SNE visualization of the learned feature by ViT-B/16.**</p>
+
+## License
+This repository is released under the Apache 2.0 license as found in the [LICENSE](LICENSE) file.
+
+## Citation
+If you find this repository useful, please consider giving a star :star: and citation :beer::
+```
+@inproceedings{mugs2022SSL,
+  title={Mugs: A Multi-Granular Self-Supervised Learning Framework},
+  author={Pan Zhou and Yichen Zhou and Chenyang Si and Weihao Yu and Teck Khim Ng and Shuicheng Yan},
+  booktitle={arXiv preprint arXiv:2203.14415},
+  year={2022}
+}

main.py

@@ -0,0 +1,814 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Mugs training code
+"""
+import argparse
+import datetime
+import json
+import math
+import os
+import sys
+import time
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+from torchvision import models as torchvision_models
+
+import utils
+from src.loss import get_multi_granular_loss
+from src.model import get_model
+from src.multicropdataset import data_prefetcher, get_dataset
+from src.optimizer import cancel_gradients_last_layer, get_optimizer, clip_gradients
+
+torchvision_archs = sorted(
+    name
+    for name in torchvision_models.__dict__
+    if name.islower()
+    and not name.startswith("__")
+    and callable(torchvision_models.__dict__[name])
+)
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser("Mugs", add_help=False)
+
+    ##======== Model parameters ============
+    parser.add_argument(
+        "--arch",
+        type=str,
+        default="vit_small",
+        choices=["vit_small", "vit_base", "vit_large"],
+        help="""Name of architecture to train.""",
+    )
+    parser.add_argument(
+        "--patch_size",
+        type=int,
+        default=16,
+        help="""Size in pixels
+        of input square patches - default 16 (for 16x16 patches). Using smaller
+        values leads to better performance but requires more memory. Applies only
+        for ViTs (vit_small and vit_base). If <16, we recommend disabling
+        mixed precision training (--use_fp16 false) to avoid unstabilities.""",
+    )
+
+    ##======== Training/Optimization parameters ============
+    parser.add_argument(
+        "--momentum_teacher",
+        type=float,
+        default=0.996,
+        help="""Base EMA
+        parameter for teacher update. The value is increased to 1 during training with
+        cosine schedule. We recommend setting a higher value with small batches: for
+        example use 0.9995 with batch size of 256.""",
+    )
+    parser.add_argument(
+        "--use_fp16",
+        type=utils.bool_flag,
+        default=False,
+        help="""Whether or not
+        to use half precision for training. Improves training time and memory requirements,
+        but can provoke instability and slight decay of performance. We recommend disabling
+        mixed precision if the loss is unstable, if reducing the patch size or if training
+        with bigger ViTs.""",
+    )
+    parser.add_argument(
+        "--weight_decay",
+        type=float,
+        default=0.04,
+        help="""Initial value of the
+        weight decay. With ViT, a smaller value at the beginning of training works well.""",
+    )
+    parser.add_argument(
+        "--weight_decay_end",
+        type=float,
+        default=0.2,
+        help="""Final value of the
+        weight decay. We use a cosine schedule for WD and using a larger decay by
+        the end of training improves performance for ViTs.""",
+    )
+    parser.add_argument(
+        "--clip_grad",
+        type=float,
+        default=3.0,
+        help="""Maximal parameter
+        gradient norm if using gradient clipping. Clipping with norm .3 ~ 1.0 can
+        help optimization for larger ViT architectures. 0 for disabling.""",
+    )
+    parser.add_argument(
+        "--batch_size_per_gpu",
+        type=int,
+        default=64,
+        help="Per-GPU batch-size : number of distinct images loaded on one GPU.",
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=100, help="Number of epochs of training."
+    )
+    parser.add_argument(
+        "--warmup_epochs",
+        default=10,
+        type=int,
+        help="""Number of epochs for the linear learning-rate warm up.=""",
+    )
+    parser.add_argument(
+        "--freeze_last_layer",
+        type=int,
+        default=1,
+        help="""Number of epochs during
+        which we keep the output layer fixed for the group supervision loss. Typically doing so during
+        the first epoch helps training. Try increasing this value if the loss does not decrease.""",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.0008,
+        help="""Learning rate at the end of
+        linear warmup (highest LR used during training). The learning rate is linearly scaled
+        with the batch size, and specified here for a reference batch size of 256.""",
+    )
+    parser.add_argument(
+        "--patch_embed_lr_mult",
+        type=float,
+        default=0.2,
+        help="""For patch
+        embedding layer, its learning rate is lr * patch_embed_lr_mult (<1.0) in most case, which
+        stables training and also slightly improve the performance.""",
+    )
+    parser.add_argument(
+        "--min_lr",
+        type=float,
+        default=1e-6,
+        help="""Target LR at the
+        end of optimization. We use a cosine LR schedule with linear warmup.""",
+    )
+    parser.add_argument(
+        "--optimizer",
+        type=str,
+        default="adamw",
+        choices=["adamw", "sgd", "lars"],
+        help="""Type of optimizer. We recommend using adamw
+        with ViTs.""",
+    )
+    parser.add_argument(
+        "--drop_path_rate", type=float, default=0.1, help="""stochastic depth rate"""
+    )
+
+    ##========  Multi-granular supervisions (instance/local-group/group supervisions) ==========
+    parser.add_argument(
+        "--loss_weights",
+        type=float,
+        nargs="+",
+        default=[1.0, 1.0, 1.0],
+        help="""three loss weights for instance, local-group, group supervision losses in turn""",
+    )
+
+    parser.add_argument(
+        "--use_bn_in_head",
+        type=utils.bool_flag,
+        default=False,
+        help="Whether to use batch normalizations in the three projection heads (Default: False)",
+    )
+    parser.add_argument(
+        "--norm_before_pred",
+        type=utils.bool_flag,
+        default=True,
+        help="""Whether to use batch normalizations after projection heads (namely before
+        prediction heads) in instance and local-group supervisions. (Default: False)""",
+    )
+
+    # parameters for instance discrimination supervision
+    parser.add_argument(
+        "--instance_out_dim",
+        type=int,
+        default=256,
+        help="""output dimention in the projection and prediction heads.""",
+    )
+    parser.add_argument(
+        "--instance_queue_size",
+        type=int,
+        default=65536,
+        help="""the queue size of the memory to store the negative keys.""",
+    )
+    parser.add_argument(
+        "--instance_temp",
+        type=float,
+        default=0.2,
+        help="""the temperature parameters for the infoNCE loss in instance supervision.""",
+    )
+
+    # parameters for local-group discrimination supervision
+    parser.add_argument(
+        "--local_group_out_dim",
+        type=int,
+        default=256,
+        help="""output dimention in the projection and prediction heads.""",
+    )
+    parser.add_argument(
+        "--local_group_knn_top_n",
+        type=int,
+        default=8,
+        help="how many neighbors we use to aggregate for a local-group",
+    )
+    parser.add_argument(
+        "--local_group_queue_size",
+        type=int,
+        default=65536,
+        help="""the queue sizes of the memory to store the negative keys for infoNCE loss and
+        another memory size to store the weak augmentated samples for local-group aggregation.""",
+    )
+    parser.add_argument(
+        "--local_group_temp",
+        type=float,
+        default=0.2,
+        help="""the temperature parameters for the infoNCE loss in instance supervision.""",
+    )
+
+    ## parameters for group discrimination supervision
+    parser.add_argument(
+        "--group_out_dim",
+        type=int,
+        default=65536,
+        help="""output dimention in the prediction heads.""",
+    )
+    parser.add_argument(
+        "--group_bottleneck_dim",
+        type=float,
+        default=256,
+        help="""head bottleneck dimention in the prediction heads.""",
+    )
+    parser.add_argument(
+        "--norm_last_layer",
+        type=utils.bool_flag,
+        default=True,
+        help="""Whether or not to weight normalize the last layer of the group supervision head.
+        Not normalizing leads to better performance but can make the training unstable. We
+        typically set this paramater to False with vit_small and True with vit_base and vit_large.""",
+    )
+
+    parser.add_argument(
+        "--group_student_temp",
+        type=float,
+        default=0.1,
+        help="""the temperature parameters for the clustering loss in student output.""",
+    )
+    parser.add_argument(
+        "--group_warmup_teacher_temp",
+        default=0.04,
+        type=float,
+        help="""Initial value for the teacher temperature: 0.04 works well in most cases.
+        Try decreasing it if the training loss does not decrease.""",
+    )
+    parser.add_argument(
+        "--group_teacher_temp",
+        default=0.04,
+        type=float,
+        help="""Final value
+        (after linear warmup) of the teacher temperature. For most experiments, anything above
+        0.07 is unstable. We recommend starting with the default value of 0.04 and increase
+        this slightly if needed.""",
+    )
+    parser.add_argument(
+        "--group_warmup_teacher_temp_epochs",
+        default=0,
+        type=int,
+        help="""Number of warmup epochs for the teacher temperature (Default: 30).""",
+    )
+
+    ##======== augmentation parameters  ============
+    # Multi-crop parameters
+    parser.add_argument(
+        "--global_crops_scale",
+        type=float,
+        nargs="+",
+        default=(0.25, 1.0),
+        help="""Scale range of the cropped image before resizing, relatively to the origin image.
+        Used for large global view cropping. When disabling multi-crop (--local_crops_number 0), we
+        recommand using a wider range of scale ("--global_crops_scale 0.14 1." for example)""",
+    )
+    parser.add_argument(
+        "--local_crops_number",
+        type=int,
+        default=10,
+        help="""Number of small
+        local views to generate. Set this parameter to 0 to disable multi-crop training.
+        When disabling multi-crop we recommend to use "--global_crops_scale 0.14 1." """,
+    )
+    parser.add_argument(
+        "--local_crops_scale",
+        type=float,
+        nargs="+",
+        default=(0.05, 0.25),
+        help="""Scale range of the cropped image before resizing, relatively to the origin image.
+        Used for small local view cropping of multi-crop.""",
+    )
+    # strong augmentation parameters
+    parser.add_argument(
+        "--timm_auto_augment_par",
+        type=str,
+        default="rand-m9-mstd0.5-inc1",
+        help="""the parameters for the AutoAugment used in DeiT.""",
+    )
+    parser.add_argument(
+        "--color_aug",
+        type=utils.bool_flag,
+        default=False,
+        help="""after AutoAugment, whether we further perform color augmentation. (Default: False).""",
+    )
+    parser.add_argument(
+        "--size_crops",
+        type=int,
+        default=[96],
+        nargs="+",
+        help="""the small crop size. Note we use multi-crop strategy, namely two 224-sized crops +
+        ten 96-sized crops. (Default: 96)""",
+    )
+    parser.add_argument(
+        "--strong_ratio",
+        type=float,
+        default=0.45,
+        help="""the ratio of image augmentation for the AutoAugment used in DeiT.""",
+    )
+    parser.add_argument(
+        "--re_prob",
+        type=float,
+        default=0.25,
+        help="""the re-prob parameter of image augmentation for the AutoAugment used in DeiT.""",
+    )
+    parser.add_argument(
+        "--vanilla_weak_augmentation",
+        type=utils.bool_flag,
+        default=False,
+        help="""Whether we use the same augmentation in DINO, namely only using weak augmentation.""",
+    )
+    parser.add_argument(
+        "--prob",
+        type=float,
+        default=0.5,
+        help="""When we use strong augmentation and weak augmentation, the ratio of images to
+        be cropped with strong augmentation.""",
+    )
+
+    ##======== Misc ============
+    parser.add_argument(
+        "--data_path",
+        default="/dataset/imageNet100_sicy/train/",
+        type=str,
+        help="""Please specify path to the ImageNet training data.""",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default="./exp/",
+        type=str,
+        help="""Path to save logs and checkpoints.""",
+    )
+    parser.add_argument(
+        "--saveckp_freq",
+        default=50,
+        type=int,
+        help="""Save checkpoint every x epochs.""",
+    )
+    parser.add_argument("--seed", default=0, type=int, help="""Random seed.""")
+    parser.add_argument(
+        "--num_workers",
+        default=12,
+        type=int,
+        help="""Number of data loading workers per GPU.""",
+    )
+    parser.add_argument(
+        "--dist_url",
+        default="env://",
+        type=str,
+        help="""url used to set up
+        distributed training; see https://pytorch.org/docs/stable/distributed.html""",
+    )
+    parser.add_argument(
+        "--local_rank",
+        default=0,
+        type=int,
+        help="""local rank for distrbuted training.""",
+    )
+    parser.add_argument(
+        "--rank", default=0, type=int, help="""rank for distrbuted training."""
+    )
+    parser.add_argument(
+        "--world_size",
+        default=1,
+        type=int,
+        help="""world size for distrbuted training.""",
+    )
+
+    parser.add_argument(
+        "--use_prefetcher",
+        type=utils.bool_flag,
+        default=True,
+        help="""whether we use prefetcher which can accerelate the training speed.""",
+    )
+    parser.add_argument(
+        "--debug",
+        type=utils.bool_flag,
+        default=False,
+        help="""whether we debug. if yes, we only load small fraction of training data to reduce data reading time.""",
+    )
+    parser.add_argument(
+        "--ddpjob",
+        default=False,
+        type=utils.bool_flag,
+        help="""whether we use ddp job. We suggest to use it for distributed training. For single GPUs
+        or Node, you can close it.""",
+    )
+
+    return parser
+
+
+def train_mugs(args):
+    """
+    main training code for Mugs, including building dataloader, models, losses, optimizers, etc
+    """
+    ##======== prepare logger for more detailed logs ============
+    logger = utils.get_logger(args.output_dir + "/train.log")
+    logger.info(args)
+    if args.output_dir and utils.is_main_process():
+        with (Path(args.output_dir) / "log.txt").open("a") as f:
+            f.write(str(args) + "\n")
+
+    ##======== initilize distribution ============
+    if args.ddpjob is True:
+        utils.init_distributed_ddpjob(args)
+    else:
+        utils.init_distributed_mode(args)
+
+    ##======== fix seed for reproduce ============
+    utils.fix_random_seeds(args.seed)
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print(
+        "\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))
+    )
+    cudnn.benchmark = True
+    cudnn.deterministic = True
+
+    ##======== get the training dataset/loader ============
+    data_loader = get_dataset(args)
+    logger.info(f"Data loaded: there are {len(data_loader.dataset)} images.")
+
+    ##====== build  student and teacher networks (vit_small, vit_base, vit_large) =========
+    student, teacher, student_mem, teacher_mem = get_model(args)
+
+    # move networks to gpu
+    student, teacher = student.cuda(), teacher.cuda()
+    student_mem, teacher_mem = student_mem.cuda(), teacher_mem.cuda()
+
+    # synchronize batch norms (if any)
+    if utils.has_batchnorms(student):
+        student = nn.SyncBatchNorm.convert_sync_batchnorm(student)
+        teacher = nn.SyncBatchNorm.convert_sync_batchnorm(teacher)
+        # we need DDP wrapper to have synchro batch norms working...
+        teacher = nn.parallel.DistributedDataParallel(teacher, device_ids=[args.gpu])
+        teacher_without_ddp = teacher.module
+    else:
+        # teacher_without_ddp and teacher are the same thing
+        teacher_without_ddp = teacher
+    student = nn.parallel.DistributedDataParallel(student, device_ids=[args.gpu])
+    # teacher and student start with the same weights
+    teacher_without_ddp.load_state_dict(student.module.state_dict(), strict=False)
+
+    # there is no backpropagation through the teacher, so no need for gradients
+    for p in teacher.parameters():
+        p.requires_grad = False
+    print(f"Student and Teacher are built: they are both {args.arch} network.")
+
+    ##======== get  multi granular losses and their loss weights ============
+    all_losses, all_weights = get_multi_granular_loss(args)
+
+    ##======== preparing optimizer ============
+    optimizer, fp16_scaler, lr_schedule, wd_schedule, momentum_schedule = get_optimizer(
+        student, len(data_loader), args
+    )
+
+    ##======== optionally resume training ============
+    to_restore = {"epoch": 0}
+    utils.restart_from_checkpoint(
+        os.path.join(args.output_dir, "checkpoint.pth"),
+        run_variables=to_restore,
+        student=student,
+        teacher=teacher,
+        optimizer=optimizer,
+        fp16_scaler=fp16_scaler,
+        student_mem=student_mem,
+        teacher_mem=teacher_mem,
+        **all_losses,
+    )
+    start_epoch = to_restore["epoch"]
+
+    ##======== Starting Mugs training ============
+    logger.info("Starting Mugs training !")
+    start_time = time.time()
+    for epoch in range(start_epoch, args.epochs):
+        t1 = time.time()
+        data_loader.sampler.set_epoch(epoch)
+
+        ##======== training one epoch of Mugs ============
+        train_stats = train_one_epoch(
+            student,
+            teacher,
+            teacher_without_ddp,
+            all_losses,
+            all_weights,
+            data_loader,
+            optimizer,
+            lr_schedule,
+            wd_schedule,
+            momentum_schedule,
+            epoch,
+            fp16_scaler,
+            student_mem,
+            teacher_mem,
+            logger,
+            args,
+        )
+
+        ##======== save model checkpoint ============
+        save_dict = {
+            "student": student.state_dict(),
+            "teacher": teacher.state_dict(),
+            "student_mem": student_mem.state_dict()
+            if student_mem is not None
+            else None,
+            "teacher_mem": teacher_mem.state_dict()
+            if teacher_mem is not None
+            else None,
+            "optimizer": optimizer.state_dict(),
+            "epoch": epoch + 1,
+            "args": args,
+        }
+        granular_loss_dicts = {}
+        for name, loss in all_losses.items():
+            granular_loss_dicts[name] = loss.state_dict()
+        save_dict.update(granular_loss_dicts)
+
+        if fp16_scaler is not None:
+            save_dict["fp16_scaler"] = fp16_scaler.state_dict()
+
+        utils.save_on_master(save_dict, os.path.join(args.output_dir, "checkpoint.pth"))
+        if args.saveckp_freq and epoch % args.saveckp_freq == 0:
+            utils.save_on_master(
+                save_dict, os.path.join(args.output_dir, f"checkpoint{epoch:04}.pth")
+            )
+
+        ##======== writing logs ============
+        log_stats = {**{f"{k}": v for k, v in train_stats.items()}, "epoch": epoch}
+        if utils.is_main_process():
+            with (Path(args.output_dir) / "log.txt").open("a") as f:
+                f.write(json.dumps(log_stats) + "\n")
+
+            t2 = time.time()
+            log_results = ""
+            for k, v in train_stats.items():
+                log_results += "%s: %.6f, " % (k, v)
+            logger.info(
+                "%d-epoch: %s remaining time %.2f hours"
+                % (epoch, log_results, (t2 - t1) * (args.epochs - epoch) / 3600.0)
+            )
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    logger.info("Training time {}".format(total_time_str))
+
+
+def train_one_epoch(
+    student,
+    teacher,
+    teacher_without_ddp,
+    all_losses,
+    all_weights,
+    data_loader,
+    optimizer,
+    lr_schedule,
+    wd_schedule,
+    momentum_schedule,
+    epoch,
+    fp16_scaler,
+    student_mem,
+    teacher_mem,
+    logger,
+    args,
+):
+    """
+    main training code for each epoch
+    """
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    prefetcher = data_prefetcher(data_loader, fp16=(fp16_scaler is not None))
+    images, weak_aug_flags = prefetcher.next()
+    epoch_it = 0
+    while images is not None:
+        #  Step 1. update weight decay and learning rate according to their schedule
+        it = len(data_loader) * epoch + epoch_it  # global training iteration
+        for _, param_group in enumerate(optimizer.param_groups):
+            lr_mult = 1.0
+            if "patch_embed" in param_group["name"]:
+                lr_mult = args.patch_embed_lr_mult
+            param_group["lr"] = lr_schedule[it] * lr_mult
+            if param_group.get("apply_wd", True):  # only the first group is regularized
+                param_group["weight_decay"] = wd_schedule[it]
+
+        granular_losses = OrderedDict()
+        total_loss = 0
+        with torch.cuda.amp.autocast(fp16_scaler is not None):
+            ## Step 2. forward images into teacher and student to obtain the
+            # features/superivisons for the three granular superivison losses
+            (
+                teacher_instance_target,
+                teacher_local_group_target,
+                teacher_group_target,
+                teacher_memory_tokens,
+            ) = teacher(
+                images[:2],
+                return_target=True,
+                local_group_memory_inputs={"mem": teacher_mem},
+            )
+
+            (
+                student_instance_target,
+                student_local_group_target,
+                student_group_target,
+                student_memory_tokens,
+            ) = student(
+                images[2:],
+                return_target=False,
+                local_group_memory_inputs={"mem": student_mem},
+            )
+
+            ## Step 3. compute the three granular supervision losses, including instance,
+            # local-group, group supervision losses
+            weigts_sum, total_loss, granular_losses = 0.0, 0.0, OrderedDict()
+            # instance loss
+            loss_cls, loss_weight = (
+                all_losses["instance-sup."],
+                all_weights["instance-sup."],
+            )
+            if loss_weight > 0:
+                instance_loss = loss_cls(
+                    student_instance_target, teacher_instance_target, epoch
+                )
+                weigts_sum, total_loss = (
+                    weigts_sum + loss_weight,
+                    total_loss + instance_loss,
+                )
+                granular_losses["instance-sup."] = instance_loss.item()
+
+            # local group loss
+            loss_cls, loss_weight = (
+                all_losses["local-group-sup."],
+                all_weights["local-group-sup."],
+            )
+            if loss_weight > 0:
+                local_group_loss = loss_cls(
+                    student_local_group_target, teacher_local_group_target, epoch
+                )
+                weigts_sum, total_loss = (
+                    weigts_sum + loss_weight,
+                    total_loss + local_group_loss,
+                )
+                granular_losses["local-group-sup."] = local_group_loss.item()
+
+            # group loss
+            loss_cls, loss_weight = all_losses["group-sup."], all_weights["group-sup."]
+            if loss_weight > 0:
+                group_loss = loss_cls(student_group_target, teacher_group_target, epoch)
+                weigts_sum, total_loss = (
+                    weigts_sum + loss_weight,
+                    total_loss + group_loss,
+                )
+                granular_losses["group-sup."] = group_loss.item()
+
+            # average loss
+            total_loss /= weigts_sum
+
+            ## ## Step 4. update the memory buffer for local-group supervision losses.
+            # for student, we only update memory by the image of size 224 and weak augmentations
+            student_features = (student_memory_tokens.chunk(2))[0]
+            len_weak = student_mem._dequeue_and_enqueue(
+                student_features, 
+                weak_aug_flags, 
+            )
+
+            teacher_weak = (teacher_memory_tokens.chunk(2))[0]
+            _ = teacher_mem._dequeue_and_enqueue(teacher_weak, None)
+
+        if not math.isfinite(total_loss.item()):
+            print("Loss is {}, stopping training".format(total_loss.item()), force=True)
+            sys.exit(1)
+
+        ## Step 5. student and teacher update
+        # student update
+        optimizer.zero_grad()
+        if fp16_scaler is None:
+            total_loss.backward()
+            if args.clip_grad:
+                clip_grad = args.clip_grad
+                if epoch > 100 and args.arch == "vit_large":
+                    clip_grad = args.clip_grad / 10.0
+                _ = clip_gradients(student, clip_grad)
+            cancel_gradients_last_layer(epoch, student, args.freeze_last_layer)
+            optimizer.step()
+        else:
+            fp16_scaler.scale(total_loss).backward()
+            if args.clip_grad:
+                clip_grad = args.clip_grad
+                if epoch > 100 and args.arch == "vit_large":
+                    clip_grad = args.clip_grad /10.0
+                fp16_scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                _ = clip_gradients(student, clip_grad)
+            cancel_gradients_last_layer(epoch, student, args.freeze_last_layer)
+            fp16_scaler.step(optimizer)
+            fp16_scaler.update()
+
+        # EMA update for the teacher
+        with torch.no_grad():
+            m = momentum_schedule[it]  # momentum parameter
+            for param_q, param_k in zip(
+                student.module.backbone.parameters(),
+                teacher_without_ddp.backbone.parameters(),
+            ):
+                param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)
+
+            if teacher_without_ddp.instance_head is not None:
+                for param_q, param_k in zip(
+                    student.module.instance_head.parameters(),
+                    teacher_without_ddp.instance_head.parameters(),
+                ):
+                    param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)
+
+            if teacher_without_ddp.local_group_head is not None:
+                for param_q, param_k in zip(
+                    student.module.local_group_head.parameters(),
+                    teacher_without_ddp.local_group_head.parameters(),
+                ):
+                    param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)
+
+            if teacher_without_ddp.group_head is not None:
+                for param_q, param_k in zip(
+                    student.module.group_head.parameters(),
+                    teacher_without_ddp.group_head.parameters(),
+                ):
+                    param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)
+
+        ## Step 6. load images
+        images, weak_aug_flags = prefetcher.next()
+        epoch_it += 1
+
+        ## Step 7. logging
+        torch.cuda.synchronize()
+        metric_logger.update(loss=total_loss.item())
+        for loss_name, loss_value in granular_losses.items():
+            metric_logger.update(**{loss_name: loss_value})
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+        metric_logger.update(wd=optimizer.param_groups[0]["weight_decay"])
+
+        if epoch_it % 500 == 0 and args.rank == 0:  # and epoch_it < 10:
+            log_results = ""
+            for _, loss_name in enumerate(all_losses):
+                if all_weights[loss_name] > 0:
+                    log_results += "%s: %.6f," % (
+                        loss_name,
+                        metric_logger.meters[loss_name].global_avg,
+                    )
+            logger.info(
+                "%d-epoch (%d/%d): total loss %.6f, %s, lr %.4e, wd %.4e, weak aug. ratio %.1f"
+                % (
+                    epoch,
+                    it,
+                    len(data_loader),
+                    metric_logger.meters["loss"].global_avg,
+                    log_results,
+                    optimizer.param_groups[0]["lr"],
+                    optimizer.param_groups[0]["weight_decay"],
+                    len_weak / len(weak_aug_flags) / args.world_size,
+                )
+            )
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("Mugs", parents=[get_args_parser()])
+    args = parser.parse_args()
+    if not os.path.exists(args.output_dir):
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+    train_mugs(args)

pretraining.sh

@@ -0,0 +1,231 @@
+#!/bin/bash
+DATASET_ROOT=/dataset/imageNet100_sicy/train/ #/raid/common/imagenet-raw/
+
+## train ViT-small for 100 epochs
+OUTPUT_ROOT=./exps/vit_small_100ep
+NPROC_PER_NODE=8 # GPU numbers
+BATCH_SIZE_PER_GPU=64
+DEBUG=false # debug = true, then we only load subset of the whole training dataset
+python -m torch.distributed.launch --nproc_per_node=$NPROC_PER_NODE main.py \
+	--data_path $DATASET_ROOT \
+	--output_dir $OUTPUT_ROOT \
+	--arch vit_small \
+	--instance_queue_size 65536 \
+	--local_group_queue_size 65536 \
+	--use_bn_in_head false \
+	--instance_out_dim 256 \
+	--instance_temp 0.2 \
+	--local_group_out_dim 256 \
+	--local_group_temp 0.2 \
+	--local_group_knn_top_n 8 \
+	--group_out_dim 65536 \
+	--group_student_temp 0.1 \
+	--group_warmup_teacher_temp 0.04 \
+	--group_teacher_temp 0.04 \
+	--group_warmup_teacher_temp_epochs 0 \
+	--norm_last_layer false \
+	--norm_before_pred true \
+	--batch_size_per_gpu $BATCH_SIZE_PER_GPU \
+	--epochs 100 \
+	--warmup_epochs 10 \
+	--clip_grad 3.0 \
+	--lr 0.0008 \
+	--min_lr 1e-06 \
+	--patch_embed_lr_mult 0.2 \
+	--drop_path_rate 0.1 \
+	--weight_decay 0.04 \
+	--weight_decay_end 0.2 \
+	--freeze_last_layer 1 \
+	--momentum_teacher 0.996 \
+	--use_fp16 false \
+	--local_crops_number 10 \
+	--size_crops 96 \
+	--global_crops_scale 0.25 1 \
+	--local_crops_scale 0.05 0.25 \
+	--timm_auto_augment_par rand-m9-mstd0.5-inc1 \
+	--prob 0.5 \
+	--use_prefetcher true \
+	--debug $DEBUG
+
+## train ViT-small for 300 epochs
+OUTPUT_ROOT=./exps/vit_small_300ep
+NPROC_PER_NODE=16 # GPU numbers
+BATCH_SIZE_PER_GPU=64
+DEBUG=false # debug = true, then we only load subset of the whole training dataset
+python -m torch.distributed.launch --nproc_per_node=$NPROC_PER_NODE main.py \
+	--data_path $DATASET_ROOT \
+	--output_dir $OUTPUT_ROOT \
+	--arch vit_small \
+	--instance_queue_size 65536 \
+	--local_group_queue_size 65536 \
+	--use_bn_in_head false \
+	--instance_out_dim 256 \
+	--instance_temp 0.2 \
+	--local_group_out_dim 256 \
+	--local_group_temp 0.2 \
+	--local_group_knn_top_n 8 \
+	--group_out_dim 65536 \
+	--group_student_temp 0.1 \
+	--group_warmup_teacher_temp 0.04 \
+	--group_teacher_temp 0.07 \
+	--group_warmup_teacher_temp_epochs 30 \
+	--norm_last_layer false \
+	--norm_before_pred true \
+	--batch_size_per_gpu $BATCH_SIZE_PER_GPU \
+	--epochs 300 \
+	--warmup_epochs 10 \
+	--clip_grad 3.0 \
+	--lr 0.0008 \
+	--min_lr 1e-06 \
+	--patch_embed_lr_mult 0.2 \
+	--drop_path_rate 0.1 \
+	--weight_decay 0.04 \
+	--weight_decay_end 0.1 \
+	--freeze_last_layer 1 \
+	--momentum_teacher 0.996 \
+	--use_fp16 false \
+	--local_crops_number 10 \
+	--size_crops 96 \
+	--global_crops_scale 0.25 1 \
+	--local_crops_scale 0.05 0.25 \
+	--timm_auto_augment_par rand-m9-mstd0.5-inc1 \
+	--prob 0.5 \
+	--use_prefetcher true \
+	--debug $DEBUG
+
+## train ViT-small for 800 epochs
+NPROC_PER_NODE=16 # GPU numbers
+BATCH_SIZE_PER_GPU=64
+DEBUG=false # debug = true, then we only load subset of the whole training dataset
+python -m torch.distributed.launch --nproc_per_node=$NPROC_PER_NODE main.py \
+	--data_path $DATASET_ROOT \
+	--output_dir $OUTPUT_ROOT \
+	--arch vit_small \
+	--instance_queue_size 65536 \
+	--local_group_queue_size 65536 \
+	--use_bn_in_head false \
+	--instance_out_dim 256 \
+	--instance_temp 0.2 \
+	--local_group_out_dim 256 \
+	--local_group_temp 0.2 \
+	--local_group_knn_top_n 8 \
+	--group_out_dim 65536 \
+	--group_student_temp 0.1 \
+	--group_warmup_teacher_temp 0.04 \
+	--group_teacher_temp 0.07 \
+	--group_warmup_teacher_temp_epochs 30 \
+	--norm_last_layer false \
+	--norm_before_pred true \
+	--batch_size_per_gpu $BATCH_SIZE_PER_GPU \
+	--epochs 800 \
+	--warmup_epochs 10 \
+	--clip_grad 3.0 \
+	--lr 0.0008 \
+	--min_lr 1e-06 \
+	--patch_embed_lr_mult 0.2 \
+	--drop_path_rate 0.1 \
+	--weight_decay 0.04 \
+	--weight_decay_end 0.1 \
+	--freeze_last_layer 1 \
+	--momentum_teacher 0.996 \
+	--use_fp16 false \
+	--local_crops_number 10 \
+	--size_crops 96 \
+	--global_crops_scale 0.25 1 \
+	--local_crops_scale 0.05 0.25 \
+	--timm_auto_augment_par rand-m9-mstd0.5-inc1 \
+	--prob 0.5 \
+	--use_prefetcher true \
+	--debug $DEBUG
+
+## train ViT-base for 400 epochs
+OUTPUT_ROOT=./exps/vit_base_400ep
+NPROC_PER_NODE=24 # GPU numbers
+BATCH_SIZE_PER_GPU=42
+DEBUG=false # debug = true, then we only load subset of the whole training dataset
+python -m torch.distributed.launch --nproc_per_node=$NPROC_PER_NODE main.py \
+	--data_path $DATASET_ROOT \
+	--output_dir $OUTPUT_ROOT \
+	--arch vit_base \
+	--instance_queue_size 65536 \
+	--local_group_queue_size 65536 \
+	--use_bn_in_head false \
+	--instance_out_dim 256 \
+	--instance_temp 0.2 \
+	--local_group_out_dim 256 \
+	--local_group_temp 0.2 \
+	--local_group_knn_top_n 8 \
+	--group_out_dim 65536 \
+	--group_student_temp 0.1 \
+	--group_warmup_teacher_temp 0.04 \
+	--group_teacher_temp 0.07 \
+	--group_warmup_teacher_temp_epochs 50 \
+	--norm_last_layer false \
+	--norm_before_pred true \
+	--batch_size_per_gpu $BATCH_SIZE_PER_GPU \
+	--epochs 400 \
+	--warmup_epochs 10 \
+	--clip_grad 3.0 \
+	--lr 0.0008 \
+	--min_lr 2e-06 \
+	--patch_embed_lr_mult 0.2 \
+	--drop_path_rate 0.1 \
+	--weight_decay 0.04 \
+	--weight_decay_end 0.1 \
+	--freeze_last_layer 3 \
+	--momentum_teacher 0.996 \
+	--use_fp16 false \
+	--local_crops_number 10 \
+	--size_crops 96 \
+	--global_crops_scale 0.25 1 \
+	--local_crops_scale 0.05 0.25 \
+	--timm_auto_augment_par rand-m9-mstd0.5-inc1 \
+	--prob 0.5 \
+	--use_prefetcher true \
+	--debug $DEBUG
+
+## train ViT-large for 250 epochs
+OUTPUT_ROOT=./exps/vit_large_250ep
+NPROC_PER_NODE=40 # GPU numbers
+BATCH_SIZE_PER_GPU=16
+DEBUG=false # debug = true, then we only load subset of the whole training dataset
+python -m torch.distributed.launch --nproc_per_node=$NPROC_PER_NODE main.py \
+	--data_path $DATASET_ROOT \
+	--output_dir $OUTPUT_ROOT \
+	--arch vit_large \
+	--instance_queue_size 65536 \
+	--local_group_queue_size 65536 \
+	--use_bn_in_head false \
+	--instance_out_dim 256 \
+	--instance_temp 0.2 \
+	--local_group_out_dim 256 \
+	--local_group_temp 0.2 \
+	--local_group_knn_top_n 8 \
+	--group_out_dim 65536 \
+	--group_student_temp 0.1 \
+	--group_warmup_teacher_temp 0.04 \
+	--group_teacher_temp 0.07 \
+	--group_warmup_teacher_temp_epochs 50 \
+	--norm_last_layer true \
+	--norm_before_pred true \
+	--batch_size_per_gpu $BATCH_SIZE_PER_GPU \
+	--epochs 250 \
+	--warmup_epochs 10 \
+	--clip_grad 3.0 \
+	--lr 0.0015 \
+	--min_lr 1.5e-4 \
+	--patch_embed_lr_mult 0.2 \
+	--drop_path_rate 0.3 \
+	--weight_decay 0.025 \
+	--weight_decay_end 0.08 \
+	--freeze_last_layer 3 \
+	--momentum_teacher 0.996 \
+	--use_fp16 false \
+	--local_crops_number 10 \
+	--size_crops 96 \
+	--global_crops_scale 0.25 1 \
+	--local_crops_scale 0.05 0.25 \
+	--timm_auto_augment_par rand-m9-mstd0.5-inc1 \
+	--prob 0.5 \
+	--use_prefetcher true \
+	--debug $DEBUG

src/RandAugment.py

@@ -0,0 +1,506 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+implment AutoAugment, RandAugment
+Adapted from https://github.com/rwightman/pytorch-image-models/blob/master/timm/data/auto_augment.py and modified for token labeling
+"""
+import math
+import random
+import re
+
+import numpy as np
+import PIL
+from PIL import Image, ImageEnhance, ImageOps
+
+_PIL_VER = tuple([int(x) for x in PIL.__version__.split(".")[:2]])
+
+_FILL = (128, 128, 128)
+
+_MAX_LEVEL = 10.0
+
+_HPARAMS_DEFAULT = dict(
+    translate_const=250,
+    img_mean=_FILL,
+)
+
+_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
+
+
+def _interpolation(kwargs):
+    interpolation = kwargs.pop("resample", Image.BILINEAR)
+    if isinstance(interpolation, (list, tuple)):
+        return random.choice(interpolation)
+    else:
+        return interpolation
+
+
+def _check_args_tf(kwargs):
+    if "fillcolor" in kwargs and _PIL_VER < (5, 0):
+        kwargs.pop("fillcolor")
+    kwargs["resample"] = _interpolation(kwargs)
+
+
+def shear_x(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)
+
+
+def shear_y(img, factor, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)
+
+
+def translate_x_rel(img, pct, **kwargs):
+    pixels = pct * img.size[0]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_rel(img, pct, **kwargs):
+    pixels = pct * img.size[1]
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def translate_x_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
+
+
+def translate_y_abs(img, pixels, **kwargs):
+    _check_args_tf(kwargs)
+    return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
+
+
+def rotate(img, degrees, **kwargs):
+    _check_args_tf(kwargs)
+    if _PIL_VER >= (5, 2):
+        return img.rotate(degrees, **kwargs)
+    elif _PIL_VER >= (5, 0):
+        w, h = img.size
+        post_trans = (0, 0)
+        rotn_center = (w / 2.0, h / 2.0)
+        angle = -math.radians(degrees)
+        matrix = [
+            round(math.cos(angle), 15),
+            round(math.sin(angle), 15),
+            0.0,
+            round(-math.sin(angle), 15),
+            round(math.cos(angle), 15),
+            0.0,
+        ]
+
+        def transform(x, y, matrix):
+            (a, b, c, d, e, f) = matrix
+            return a * x + b * y + c, d * x + e * y + f
+
+        matrix[2], matrix[5] = transform(
+            -rotn_center[0] - post_trans[0], -rotn_center[1] - post_trans[1], matrix
+        )
+        matrix[2] += rotn_center[0]
+        matrix[5] += rotn_center[1]
+        return img.transform(img.size, Image.AFFINE, matrix, **kwargs)
+    else:
+        return img.rotate(degrees, resample=kwargs["resample"])
+
+
+def auto_contrast(img, **__):
+    return ImageOps.autocontrast(img)
+
+
+def invert(img, **__):
+    return ImageOps.invert(img)
+
+
+def equalize(img, **__):
+    return ImageOps.equalize(img)
+
+
+def solarize(img, thresh, **__):
+    return ImageOps.solarize(img, thresh)
+
+
+def solarize_add(img, add, thresh=128, **__):
+    lut = []
+    for i in range(256):
+        if i < thresh:
+            lut.append(min(255, i + add))
+        else:
+            lut.append(i)
+    if img.mode in ("L", "RGB"):
+        if img.mode == "RGB" and len(lut) == 256:
+            lut = lut + lut + lut
+        return img.point(lut)
+    else:
+        return img
+
+
+def posterize(img, bits_to_keep, **__):
+    if bits_to_keep >= 8:
+        return img
+    return ImageOps.posterize(img, bits_to_keep)
+
+
+def contrast(img, factor, **__):
+    return ImageEnhance.Contrast(img).enhance(factor)
+
+
+def color(img, factor, **__):
+    return ImageEnhance.Color(img).enhance(factor)
+
+
+def brightness(img, factor, **__):
+    return ImageEnhance.Brightness(img).enhance(factor)
+
+
+def sharpness(img, factor, **__):
+    return ImageEnhance.Sharpness(img).enhance(factor)
+
+
+def _randomly_negate(v):
+    """With 50% prob, negate the value"""
+    return -v if random.random() > 0.5 else v
+
+
+def _rotate_level_to_arg(level, _hparams):
+    # range [-30, 30]
+    level = (level / _MAX_LEVEL) * 30.0
+    level = _randomly_negate(level)
+    return (level,)
+
+
+def _enhance_level_to_arg(level, _hparams):
+    # range [0.1, 1.9]
+    return ((level / _MAX_LEVEL) * 1.8 + 0.1,)
+
+
+def _enhance_increasing_level_to_arg(level, _hparams):
+    # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend
+    # range [0.1, 1.9]
+    level = (level / _MAX_LEVEL) * 0.9
+    level = 1.0 + _randomly_negate(level)
+    return (level,)
+
+
+def _shear_level_to_arg(level, _hparams):
+    # range [-0.3, 0.3]
+    level = (level / _MAX_LEVEL) * 0.3
+    level = _randomly_negate(level)
+    return (level,)
+
+
+def _translate_abs_level_to_arg(level, hparams):
+    translate_const = hparams["translate_const"]
+    level = (level / _MAX_LEVEL) * float(translate_const)
+    level = _randomly_negate(level)
+    return (level,)
+
+
+def _translate_rel_level_to_arg(level, hparams):
+    # default range [-0.45, 0.45]
+    translate_pct = hparams.get("translate_pct", 0.45)
+    level = (level / _MAX_LEVEL) * translate_pct
+    level = _randomly_negate(level)
+    return (level,)
+
+
+def _posterize_level_to_arg(level, _hparams):
+    # As per Tensorflow TPU EfficientNet impl
+    # range [0, 4], 'keep 0 up to 4 MSB of original image'
+    # intensity/severity of augmentation decreases with level
+    return (int((level / _MAX_LEVEL) * 4),)
+
+
+def _posterize_increasing_level_to_arg(level, hparams):
+    # As per Tensorflow models research and UDA impl
+    # range [4, 0], 'keep 4 down to 0 MSB of original image',
+    # intensity/severity of augmentation increases with level
+    return (4 - _posterize_level_to_arg(level, hparams)[0],)
+
+
+def _posterize_original_level_to_arg(level, _hparams):
+    # As per original AutoAugment paper description
+    # range [4, 8], 'keep 4 up to 8 MSB of image'
+    # intensity/severity of augmentation decreases with level
+    return (int((level / _MAX_LEVEL) * 4) + 4,)
+
+
+def _solarize_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation decreases with level
+    return (int((level / _MAX_LEVEL) * 256),)
+
+
+def _solarize_increasing_level_to_arg(level, _hparams):
+    # range [0, 256]
+    # intensity/severity of augmentation increases with level
+    return (256 - _solarize_level_to_arg(level, _hparams)[0],)
+
+
+def _solarize_add_level_to_arg(level, _hparams):
+    # range [0, 110]
+    return (int((level / _MAX_LEVEL) * 110),)
+
+
+LEVEL_TO_ARG = {
+    "AutoContrast": None,
+    "Equalize": None,
+    "Invert": None,
+    "Rotate": _rotate_level_to_arg,
+    # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers
+    "Posterize": _posterize_level_to_arg,
+    "PosterizeIncreasing": _posterize_increasing_level_to_arg,
+    "PosterizeOriginal": _posterize_original_level_to_arg,
+    "Solarize": _solarize_level_to_arg,
+    "SolarizeIncreasing": _solarize_increasing_level_to_arg,
+    "SolarizeAdd": _solarize_add_level_to_arg,
+    "Color": _enhance_level_to_arg,
+    "ColorIncreasing": _enhance_increasing_level_to_arg,
+    "Contrast": _enhance_level_to_arg,
+    "ContrastIncreasing": _enhance_increasing_level_to_arg,
+    "Brightness": _enhance_level_to_arg,
+    "BrightnessIncreasing": _enhance_increasing_level_to_arg,
+    "Sharpness": _enhance_level_to_arg,
+    "SharpnessIncreasing": _enhance_increasing_level_to_arg,
+    "ShearX": _shear_level_to_arg,
+    "ShearY": _shear_level_to_arg,
+    "TranslateX": _translate_abs_level_to_arg,
+    "TranslateY": _translate_abs_level_to_arg,
+    "TranslateXRel": _translate_rel_level_to_arg,
+    "TranslateYRel": _translate_rel_level_to_arg,
+}
+
+
+NAME_TO_OP = {
+    "AutoContrast": auto_contrast,
+    "Equalize": equalize,
+    "Invert": invert,
+    "Rotate": rotate,
+    "Posterize": posterize,
+    "PosterizeIncreasing": posterize,
+    "PosterizeOriginal": posterize,
+    "Solarize": solarize,
+    "SolarizeIncreasing": solarize,
+    "SolarizeAdd": solarize_add,
+    "Color": color,
+    "ColorIncreasing": color,
+    "Contrast": contrast,
+    "ContrastIncreasing": contrast,
+    "Brightness": brightness,
+    "BrightnessIncreasing": brightness,
+    "Sharpness": sharpness,
+    "SharpnessIncreasing": sharpness,
+    "ShearX": shear_x,
+    "ShearY": shear_y,
+    "TranslateX": translate_x_abs,
+    "TranslateY": translate_y_abs,
+    "TranslateXRel": translate_x_rel,
+    "TranslateYRel": translate_y_rel,
+}
+
+_RAND_TRANSFORMS = [
+    "AutoContrast",
+    "Equalize",
+    "Invert",
+    "Rotate",
+    "Posterize",
+    "Solarize",
+    "SolarizeAdd",
+    "Color",
+    "Contrast",
+    "Brightness",
+    "Sharpness",
+    "ShearX",
+    "ShearY",
+    "TranslateXRel",
+    "TranslateYRel",
+    #'Cutout'
+]
+
+
+_RAND_INCREASING_TRANSFORMS = [
+    "AutoContrast",
+    "Equalize",
+    "Invert",
+    "Rotate",
+    "PosterizeIncreasing",
+    "SolarizeIncreasing",
+    "SolarizeAdd",
+    "ColorIncreasing",
+    "ContrastIncreasing",
+    "BrightnessIncreasing",
+    "SharpnessIncreasing",
+    "ShearX",
+    "ShearY",
+    "TranslateXRel",
+    "TranslateYRel",
+    #'Cutout'
+]
+
+
+# These experimental weights are based loosely on the relative improvements mentioned in paper.
+# They may not result in increased performance, but could likely be tuned to so.
+_RAND_CHOICE_WEIGHTS_0 = {
+    "Rotate": 0.3,
+    "ShearX": 0.2,
+    "ShearY": 0.2,
+    "TranslateXRel": 0.1,
+    "TranslateYRel": 0.1,
+    "Color": 0.025,
+    "Sharpness": 0.025,
+    "AutoContrast": 0.025,
+    "Solarize": 0.005,
+    "SolarizeAdd": 0.005,
+    "Contrast": 0.005,
+    "Brightness": 0.005,
+    "Equalize": 0.005,
+    "Posterize": 0,
+    "Invert": 0,
+}
+
+
+def _select_rand_weights(weight_idx=0, transforms=None):
+    transforms = transforms or _RAND_TRANSFORMS
+    assert weight_idx == 0  # only one set of weights currently
+    rand_weights = _RAND_CHOICE_WEIGHTS_0
+    probs = [rand_weights[k] for k in transforms]
+    probs /= np.sum(probs)
+    return probs
+
+
+class AugmentOp:
+    def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
+        hparams = hparams or _HPARAMS_DEFAULT
+        self.name = name
+        self.aug_fn = NAME_TO_OP[name]
+        self.level_fn = LEVEL_TO_ARG[name]
+        self.prob = prob
+        self.magnitude = magnitude
+        self.hparams = hparams.copy()
+        self.kwargs = dict(
+            fillcolor=hparams["img_mean"] if "img_mean" in hparams else _FILL,
+            resample=hparams["interpolation"]
+            if "interpolation" in hparams
+            else _RANDOM_INTERPOLATION,
+        )
+
+        # If magnitude_std is > 0, we introduce some randomness
+        # in the usually fixed policy and sample magnitude from a normal distribution
+        # with mean `magnitude` and std-dev of `magnitude_std`.
+        # NOTE This is my own hack, being tested, not in papers or reference impls.
+        self.magnitude_std = self.hparams.get("magnitude_std", 0)
+
+    def __call__(self, img):
+        if self.prob < 1.0 and random.random() > self.prob:
+            return img
+        magnitude = self.magnitude
+        if self.magnitude_std and self.magnitude_std > 0:
+            magnitude = random.gauss(magnitude, self.magnitude_std)
+        magnitude = min(_MAX_LEVEL, max(0, magnitude))  # clip to valid range
+        level_args = (
+            self.level_fn(magnitude, self.hparams)
+            if self.level_fn is not None
+            else tuple()
+        )
+        imgs = self.aug_fn(img, *level_args, **self.kwargs)
+
+        return imgs
+
+
+def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
+    hparams = hparams or _HPARAMS_DEFAULT
+    transforms = transforms or _RAND_TRANSFORMS
+    return [
+        AugmentOp(name, prob=0.5, magnitude=magnitude, hparams=hparams)
+        for name in transforms
+    ]
+
+
+class RandAugment:
+    """
+    Apply RandAug on image
+    """
+
+    def __init__(self, ops, num_layers=2, choice_weights=None):
+        self.ops = ops
+        self.num_layers = num_layers
+        self.choice_weights = choice_weights
+
+    def __call__(self, img):
+        # no replacement when using weighted choice
+        ops = np.random.choice(
+            self.ops, self.num_layers, replace=False, p=self.choice_weights
+        )
+        for op in ops:
+            img = op(img)
+
+        return img
+
+
+def rand_augment_transform(config_str, hparams):
+    """
+    Create a RandAugment transform
+    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
+    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
+    sections, not order sepecific determine
+        'm' - integer magnitude of rand augment
+        'n' - integer num layers (number of transform ops selected per image)
+        'w' - integer probabiliy weight index (index of a set of weights to influence choice of op)
+        'mstd' -  float std deviation of magnitude noise applied
+        'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0)
+    Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
+    'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2
+
+    :param hparams: Other hparams (kwargs) for the RandAugmentation scheme
+
+    :return: A PyTorch compatible Transform
+    """
+    magnitude = _MAX_LEVEL  # default to _MAX_LEVEL for magnitude (currently 10)
+    num_layers = 2  # default to 2 ops per image
+    weight_idx = None  # default to no probability weights for op choice
+    transforms = _RAND_TRANSFORMS
+    config = config_str.split("-")
+    assert config[0] == "rand"
+    config = config[1:]
+    for c in config:
+        cs = re.split(r"(\d.*)", c)
+        if len(cs) < 2:
+            continue
+        key, val = cs[:2]
+        if key == "mstd":
+            # noise param injected via hparams for now
+            hparams.setdefault("magnitude_std", float(val))
+        elif key == "inc":
+            if bool(val):  # this path
+                transforms = _RAND_INCREASING_TRANSFORMS
+        elif key == "m":
+            magnitude = int(val)
+        elif key == "n":
+            num_layers = int(val)
+        elif key == "w":
+            weight_idx = int(val)
+        else:
+            assert False, "Unknown RandAugment config section"
+    # magnitude 9
+    # hparams {'translate_const': 100, 'img_mean': (124, 116, 104), 'magnitude_std': 0.5}
+    # transforms ['AutoContrast', 'Equalize', 'Invert', 'Rotate', 'PosterizeIncreasing', \
+    # 'SolarizeIncreasing', 'SolarizeAdd', 'ColorIncreasing', 'ContrastIncreasing', \
+    # 'BrightnessIncreasing', 'SharpnessIncreasing', 'ShearX', 'ShearY', 'TranslateXRel', 'TranslateYRel']
+    ra_ops = rand_augment_ops(
+        magnitude=magnitude, hparams=hparams, transforms=transforms
+    )
+    choice_weights = (
+        None if weight_idx is None else _select_rand_weights(weight_idx)
+    )  ## None
+    return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)

src/dataset.py

@@ -0,0 +1,367 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+ImageFold function.
+
+Mostly copy-paste from torchvision references
+"""
+import os
+import os.path
+from typing import Any, Callable, Dict, List, Optional, Tuple, cast
+
+from PIL import Image
+from torchvision.datasets.vision import VisionDataset
+
+
+def has_file_allowed_extension(filename: str, extensions: Tuple[str, ...]) -> bool:
+    """Checks if a file is an allowed extension.
+
+    Args:
+        filename (string): path to a file
+        extensions (tuple of strings): extensions to consider (lowercase)
+
+    Returns:
+        bool: True if the filename ends with one of given extensions
+    """
+    return filename.lower().endswith(extensions)
+
+
+def is_image_file(filename: str) -> bool:
+    """Checks if a file is an allowed image extension.
+
+    Args:
+        filename (string): path to a file
+
+    Returns:
+        bool: True if the filename ends with a known image extension
+    """
+    return has_file_allowed_extension(filename, IMG_EXTENSIONS)
+
+
+def find_classes(directory: str, class_num: int) -> Tuple[List[str], Dict[str, int]]:
+    """Finds the class folders in a dataset.
+
+    See :class:`DatasetFolder` for details.
+    """
+    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
+    if not classes:
+        raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")
+    classes = classes[:class_num]
+    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
+    return classes, class_to_idx
+
+
+def make_dataset(
+    directory: str,
+    class_to_idx: Optional[Dict[str, int]] = None,
+    extensions: Optional[Tuple[str, ...]] = None,
+    is_valid_file: Optional[Callable[[str], bool]] = None,
+    class_num=10,
+) -> List[Tuple[str, int]]:
+    """Generates a list of samples of a form (path_to_sample, class).
+
+    See :class:`DatasetFolder` for details.
+
+    Note: The class_to_idx parameter is here optional and will use the logic of the ``find_classes`` function
+    by default.
+    """
+    directory = os.path.expanduser(directory)
+
+    if class_to_idx is None:
+        _, class_to_idx = find_classes(directory, class_num)
+    elif not class_to_idx:
+        raise ValueError(
+            "'class_to_index' must have at least one entry to collect any samples."
+        )
+
+    both_none = extensions is None and is_valid_file is None
+    both_something = extensions is not None and is_valid_file is not None
+    if both_none or both_something:
+        raise ValueError(
+            "Both extensions and is_valid_file cannot be None or not None at the same time"
+        )
+
+    if extensions is not None:
+
+        def is_valid_file(x: str) -> bool:
+            return has_file_allowed_extension(x, cast(Tuple[str, ...], extensions))
+
+    is_valid_file = cast(Callable[[str], bool], is_valid_file)
+
+    instances = []
+    available_classes = set()
+    for target_class in sorted(class_to_idx.keys()):
+        class_index = class_to_idx[target_class]
+        target_dir = os.path.join(directory, target_class)
+        if not os.path.isdir(target_dir):
+            continue
+        for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
+            for fname in sorted(fnames):
+                path = os.path.join(root, fname)
+                if is_valid_file(path):
+                    item = path, class_index
+                    instances.append(item)
+
+                    if target_class not in available_classes:
+                        available_classes.add(target_class)
+
+    empty_classes = set(class_to_idx.keys()) - available_classes
+    if empty_classes:
+        msg = (
+            f"Found no valid file for the classes {', '.join(sorted(empty_classes))}. "
+        )
+        if extensions is not None:
+            msg += f"Supported extensions are: {', '.join(extensions)}"
+        raise FileNotFoundError(msg)
+
+    return instances
+
+
+class DatasetFolder(VisionDataset):
+    """A generic data loader.
+
+    This default directory structure can be customized by overriding the
+    :meth:`find_classes` method.
+
+    Args:
+        root (string): Root directory path.
+        loader (callable): A function to load a sample given its path.
+        extensions (tuple[string]): A list of allowed extensions.
+            both extensions and is_valid_file should not be passed.
+        transform (callable, optional): A function/transform that takes in
+            a sample and returns a transformed version.
+            E.g, ``transforms.RandomCrop`` for images.
+        target_transform (callable, optional): A function/transform that takes
+            in the target and transforms it.
+        is_valid_file (callable, optional): A function that takes path of a file
+            and check if the file is a valid file (used to check of corrupt files)
+            both extensions and is_valid_file should not be passed.
+        class_num: how many classes will be loaded
+     Attributes:
+        classes (list): List of the class names sorted alphabetically.
+        class_to_idx (dict): Dict with items (class_name, class_index).
+        samples (list): List of (sample path, class_index) tuples
+        targets (list): The class_index value for each image in the dataset
+    """
+
+    def __init__(
+        self,
+        root: str,
+        loader: Callable[[str], Any],
+        extensions: Optional[Tuple[str, ...]] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        is_valid_file: Optional[Callable[[str], bool]] = None,
+        class_num=10,
+    ) -> None:
+        super(DatasetFolder, self).__init__(
+            root, transform=transform, target_transform=target_transform
+        )
+        classes, class_to_idx = self.find_classes(self.root, class_num=class_num)
+        samples = self.make_dataset(
+            self.root, class_to_idx, extensions, is_valid_file, class_num=class_num
+        )
+
+        self.loader = loader
+        self.extensions = extensions
+
+        self.classes = classes
+        self.class_to_idx = class_to_idx
+        self.samples = samples
+        self.targets = [s[1] for s in samples]
+
+    @staticmethod
+    def make_dataset(
+        directory: str,
+        class_to_idx: Dict[str, int],
+        extensions: Optional[Tuple[str, ...]] = None,
+        is_valid_file: Optional[Callable[[str], bool]] = None,
+        class_num=10,
+    ) -> List[Tuple[str, int]]:
+        """Generates a list of samples of a form (path_to_sample, class).
+
+        This can be overridden to e.g. read files from a compressed zip file instead of from the disk.
+
+        Args:
+            directory (str): root dataset directory, corresponding to ``self.root``.
+            class_to_idx (Dict[str, int]): Dictionary mapping class name to class index.
+            extensions (optional): A list of allowed extensions.
+                Either extensions or is_valid_file should be passed. Defaults to None.
+            is_valid_file (optional): A function that takes path of a file
+                and checks if the file is a valid file
+                (used to check of corrupt files) both extensions and
+                is_valid_file should not be passed. Defaults to None.
+            class_num: how many classes will be loaded
+        Raises:
+            ValueError: In case ``class_to_idx`` is empty.
+            ValueError: In case ``extensions`` and ``is_valid_file`` are None or both are not None.
+            FileNotFoundError: In case no valid file was found for any class.
+
+        Returns:
+            List[Tuple[str, int]]: samples of a form (path_to_sample, class)
+        """
+        if class_to_idx is None:
+            # prevent potential bug since make_dataset() would use the class_to_idx logic of the
+            # find_classes() function, instead of using that of the find_classes() method, which
+            # is potentially overridden and thus could have a different logic.
+            raise ValueError("The class_to_idx parameter cannot be None.")
+        return make_dataset(
+            directory,
+            class_to_idx,
+            extensions=extensions,
+            is_valid_file=is_valid_file,
+            class_num=class_num,
+        )
+
+    def find_classes(
+        self, directory: str, class_num: int
+    ) -> Tuple[List[str], Dict[str, int]]:
+        """Find the class folders in a dataset structured as follows::
+
+            directory/
+            ├── class_x
+            │   ├── xxx.ext
+            │   ├── xxy.ext
+            │   └── ...
+            │       └── xxz.ext
+            └── class_y
+                ├── 123.ext
+                ├── nsdf3.ext
+                └── ...
+                └── asd932_.ext
+
+        This method can be overridden to only consider
+        a subset of classes, or to adapt to a different dataset directory structure.
+
+        Args:
+            directory(str): Root directory path, corresponding to ``self.root``
+
+        Raises:
+            FileNotFoundError: If ``dir`` has no class folders.
+
+        Returns:
+            (Tuple[List[str], Dict[str, int]]): List of all classes and dictionary mapping each class to an index.
+        """
+        return find_classes(directory, class_num=class_num)
+
+    def __getitem__(self, index: int) -> Tuple[Any, Any]:
+        """
+        Args:
+            index (int): Index
+
+        Returns:
+            tuple: (sample, target) where target is class_index of the target class.
+        """
+        path, target = self.samples[index]
+        sample = self.loader(path)
+        if self.transform is not None:
+            sample = self.transform(sample)
+        # if self.target_transform is not None:
+        #     target = self.target_transform(target)
+
+        return sample  # , target
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+
+IMG_EXTENSIONS = (
+    ".jpg",
+    ".jpeg",
+    ".png",
+    ".ppm",
+    ".bmp",
+    ".pgm",
+    ".tif",
+    ".tiff",
+    ".webp",
+)
+
+
+def pil_loader(path: str) -> Image.Image:
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+    with open(path, "rb") as f:
+        img = Image.open(f)
+        return img.convert("RGB")
+
+
+# TODO: specify the return type
+def accimage_loader(path: str) -> Any:
+    import accimage
+
+    try:
+        return accimage.Image(path)
+    except IOError:
+        # Potentially a decoding problem, fall back to PIL.Image
+        return pil_loader(path)
+
+
+def default_loader(path: str) -> Any:
+    from torchvision import get_image_backend
+
+    if get_image_backend() == "accimage":
+        return accimage_loader(path)
+    else:
+        return pil_loader(path)
+
+
+class ImageFolder(DatasetFolder):
+    """A generic data loader where the images are arranged in this way by default: ::
+
+        root/dog/xxx.png
+        root/dog/xxy.png
+        root/dog/[...]/xxz.png
+
+        root/cat/123.png
+        root/cat/nsdf3.png
+        root/cat/[...]/asd932_.png
+
+    This class inherits from :class:`~torchvision.datasets.DatasetFolder` so
+    the same methods can be overridden to customize the dataset.
+
+    Args:
+        root (string): Root directory path.
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        loader (callable, optional): A function to load an image given its path.
+        is_valid_file (callable, optional): A function that takes path of an Image file
+            and check if the file is a valid file (used to check of corrupt files)
+        class_num: how many classes will be loaded
+     Attributes:
+        classes (list): List of the class names sorted alphabetically.
+        class_to_idx (dict): Dict with items (class_name, class_index).
+        imgs (list): List of (image path, class_index) tuples
+    """
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        loader: Callable[[str], Any] = default_loader,
+        is_valid_file: Optional[Callable[[str], bool]] = None,
+        class_num=10,
+    ):
+        super(ImageFolder, self).__init__(
+            root,
+            loader,
+            IMG_EXTENSIONS if is_valid_file is None else None,
+            transform=transform,
+            target_transform=target_transform,
+            is_valid_file=is_valid_file,
+            class_num=class_num,
+        )
+        self.imgs = self.samples

src/loss.py

@@ -0,0 +1,244 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+functions for building multi-granular losses.
+"""
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+from utils import concat_all_gather
+
+
+class InfoNCELoss(nn.Module):
+    """
+    vanilla infoNCEloss.
+    --ncrops: how many crops are used in student networks
+    --dim: feature dimension in queue determinted by output dimention of student network
+    --queue_size: queue size
+    --temperature: temperature parameter for infoNCEloss
+    """
+
+    def __init__(self, ncrops, dim=256, queue_size=65536, temperature=0.2):
+        super().__init__()
+        self.queue_size = queue_size
+        self.temperature = temperature
+
+        self.register_buffer("queue", torch.randn(dim, queue_size))
+        self.queue = nn.functional.normalize(self.queue, dim=0)
+
+        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
+        self.CrossEntropyLoss = nn.CrossEntropyLoss()
+        self.ncrops = ncrops
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self, keys):
+        """
+        queue update
+        """
+        keys = concat_all_gather(keys)
+        batch_size = keys.shape[0]
+        ptr = int(self.queue_ptr)
+        # replace the keys at ptr (dequeue and enqueue)
+        if ptr + batch_size <= self.queue_size:
+            self.queue[:, ptr : ptr + batch_size] = keys.T
+            ptr = (ptr + batch_size) % self.queue_size
+        else:
+            keys_t = keys.T
+            queue_remaining_size = self.queue_size - ptr
+            self.queue[:, ptr:] = keys_t[:, :queue_remaining_size]
+            self.queue[:, : batch_size - queue_remaining_size] = keys_t[
+                :, queue_remaining_size:
+            ]
+
+            ptr = batch_size - queue_remaining_size  # move pointer
+
+        self.queue_ptr[0] = ptr
+
+    # student_output, teacher_output
+    def forward(self, student_output, teacher_output, epoch):
+        """
+        Cross-entropy between softmax outputs of the teacher and student networks.
+        """
+        preds = student_output.chunk(self.ncrops)
+        targets = teacher_output.detach().chunk(2)
+        small_crop_loss, large_crop_loss = 0, 0
+        small_loss_terms, large_loss_terms = 0, 0
+        queue_feat = self.queue.clone().detach()
+
+        for t_idx, targ in enumerate(targets):
+            for p_idx, pred in enumerate(preds):
+                if t_idx == p_idx:
+                    continue
+                # positive logits: Nx1
+                l_pos = torch.einsum("nc,nc->n", [pred, targ]).unsqueeze(-1)
+                # negative logits: NxK
+                l_neg = torch.einsum("nc,ck->nk", [pred, queue_feat])
+                # logits: Nx(1+K)
+                logits = torch.cat([l_pos, l_neg], dim=1)
+                # apply temperature
+                logits /= self.temperature
+                # labels: positive key indicators
+                labels = torch.zeros(logits.shape[0], dtype=torch.long).to(
+                    logits.device
+                )
+                loss = self.CrossEntropyLoss(logits, labels)
+                if p_idx < 2:  ## large crop loss, namely loss on 224-sized images
+                    large_crop_loss += loss
+                    large_loss_terms += 1
+                else:  ## small crop loss, namely loss on 96-sized images
+                    small_crop_loss += loss
+                    small_loss_terms += 1
+            # dequeue and enqueue
+            self._dequeue_and_enqueue(targ)
+
+        large_crop_loss /= large_loss_terms
+        small_crop_loss /= small_loss_terms
+        loss = 0.5 * (large_crop_loss + small_crop_loss)
+        return loss
+
+
+class ClusteringLoss(nn.Module):
+    """
+    Clustering loss which is very simialr to the one in DINO
+    --out_dim: center dimension determinted by output dimention of student network
+    --ncrops: how many crops are used in student networks
+    --warmup_teacher_temp: Initial value for the teacher temperature
+    --teacher_temp: Final value (after linear warmup) of the teacher temperature
+    --warmup_teacher_temp_epochs: Number of warmup epochs for the teacher temperature
+    --nepochs: total training epoch
+    --student_temp: temperature parameter in student output
+    --center_momentum:  EMA parameter for center update
+    """
+
+    def __init__(
+        self,
+        out_dim,
+        ncrops,
+        warmup_teacher_temp,
+        teacher_temp,
+        warmup_teacher_temp_epochs,
+        nepochs,
+        student_temp=0.1,
+        center_momentum=0.9,
+    ):
+        super().__init__()
+        self.student_temp = student_temp
+        self.center_momentum = center_momentum
+        self.ncrops = ncrops
+        self.register_buffer("center", torch.zeros(1, out_dim))
+        # we apply a warm up for the teacher temperature because
+        # a too high temperature makes the training instable at the beginning
+        self.teacher_temp_schedule = np.concatenate(
+            (
+                np.linspace(
+                    warmup_teacher_temp, teacher_temp, warmup_teacher_temp_epochs
+                ),
+                np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp,
+            )
+        )
+
+    def forward(self, student_output, teacher_output, epoch):
+        """
+        Cross-entropy between softmax outputs of the teacher and student networks.
+        """
+        student_out = student_output / self.student_temp
+        student_out = student_out.chunk(self.ncrops)
+
+        # teacher centering and sharpening
+        temp = self.teacher_temp_schedule[epoch]
+        teacher_out = F.softmax((teacher_output - self.center) / temp, dim=-1)
+        teacher_out = teacher_out.detach().chunk(2)
+
+        loss_large_crop, loss_small_crop = 0.0, 0.0
+        loss_terms_large_crop, loss_terms_small_crop = 0, 0
+        for iq, q in enumerate(teacher_out):
+            for v in range(len(student_out)):
+                if v == iq:
+                    # we skip cases where student and teacher operate on the same view
+                    continue
+                loss = torch.sum(
+                    -q * F.log_softmax(student_out[v], dim=-1), dim=-1
+                ).mean()
+                if v < 2:
+                    loss_large_crop += loss
+                    loss_terms_large_crop += 1
+                else:
+                    loss_small_crop += loss
+                    loss_terms_small_crop += 1
+
+        self.update_center(teacher_output)
+        loss_large_crop /= loss_terms_large_crop
+        loss_small_crop /= loss_terms_small_crop
+        total_loss = 0.5 * (loss_large_crop + loss_small_crop)
+        return total_loss
+
+    @torch.no_grad()
+    def update_center(self, teacher_output):
+        """
+        Update center used for teacher output.
+        """
+        batch_center = torch.mean(teacher_output, dim=0, keepdim=False)
+        dist.all_reduce(batch_center)
+        batch_center = batch_center / dist.get_world_size()
+
+        # ema update
+        self.center = self.center * self.center_momentum + batch_center * (
+            1 - self.center_momentum
+        )
+
+
+def get_multi_granular_loss(args):
+    """
+    build the multi-granular loss
+    """
+    all_losses, all_weights = {}, {}
+
+    ## build the instance discrimination loss
+    instance_supervision_loss = InfoNCELoss(
+        args.local_crops_number + 2,
+        dim=args.instance_out_dim,
+        queue_size=args.instance_queue_size,
+        temperature=args.instance_temp,
+    ).cuda()
+    all_losses["instance-sup."] = instance_supervision_loss
+    all_weights["instance-sup."] = args.loss_weights[0]
+
+    ## build the local group discrimination loss
+    local_group_supervision = InfoNCELoss(
+        args.local_crops_number + 2,
+        dim=args.local_group_out_dim,
+        queue_size=args.local_group_queue_size,
+        temperature=args.local_group_temp,
+    ).cuda()
+    all_losses["local-group-sup."] = local_group_supervision
+    all_weights["local-group-sup."] = args.loss_weights[1]
+
+    ## build the group discrimination loss
+    group_loss = ClusteringLoss(
+        args.group_out_dim,
+        args.local_crops_number
+        + 2,  # total number of crops = 2 global crops + local_crops_number
+        args.group_warmup_teacher_temp,
+        args.group_teacher_temp,
+        args.group_warmup_teacher_temp_epochs,
+        args.epochs,
+        student_temp=args.group_student_temp,
+        center_momentum=0.9,
+    ).cuda()
+    all_losses["group-sup."] = group_loss
+    all_weights["group-sup."] = args.loss_weights[2]
+    return all_losses, all_weights

src/model.py

@@ -0,0 +1,607 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+models and functions for building student and teacher networks for multi-granular losses.
+"""
+import torch
+import torch.nn as nn
+
+import src.vision_transformer as vits
+from src.vision_transformer import trunc_normal_
+
+
+class Instance_Superivsion_Head(nn.Module):
+    """
+    a class to implement Instance Superivsion Head
+    --in_dim: input dimension of projection head
+    --hidden_dim: hidden dimension of projection head
+    --out_dim: ouput dimension of projection and prediction heads
+    --pred_hidden_dim: hidden dimension of prediction head
+    --nlayers: layer number of projection head. prediction head has nlayers-1 layer
+    --proj_bn: whether we use batch normalization in projection head
+    --pred_bn: whether we use batch normalization in prediction head
+    --norm_before_pred:  whether we use normalization before prediction head
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        hidden_dim=2048,
+        out_dim=256,
+        pred_hidden_dim=4096,
+        nlayers=3,
+        proj_bn=False,
+        pred_bn=False,
+        norm_before_pred=True,
+    ):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        self.norm_before_pred = norm_before_pred
+
+        self.projector = self._build_mlp(
+            nlayers, in_dim, hidden_dim, out_dim, use_bn=proj_bn
+        )
+
+        self.apply(self._init_weights)
+
+        self.predictor = None
+        if pred_hidden_dim > 0:  # teacher no, student yes
+            self.predictor = self._build_mlp(
+                nlayers - 1, out_dim, pred_hidden_dim, out_dim, use_bn=pred_bn
+            )
+
+    def _init_weights(self, m):
+        """
+        initilize the parameters in network
+        """
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def _build_mlp(self, num_layers, input_dim, hidden_dim, output_dim, use_bn=False):
+        """
+        build a mlp
+        """
+        mlp = []
+        for layer in range(num_layers):
+            dim1 = input_dim if layer == 0 else hidden_dim
+            dim2 = output_dim if layer == num_layers - 1 else hidden_dim
+
+            mlp.append(nn.Linear(dim1, dim2, bias=False))
+
+            if layer < num_layers - 1:
+                if use_bn:
+                    mlp.append(nn.BatchNorm1d(dim2))
+                mlp.append(nn.GELU())
+
+        return nn.Sequential(*mlp)
+
+    def forward(self, x, return_target=False):
+        """
+        forward the input through projection head for teacher and
+        projection/prediction heads for student
+        """
+        feat = self.projector(x)
+
+        if return_target:
+            feat = nn.functional.normalize(feat, dim=-1, p=2)
+            return feat
+        ## return prediction
+        if self.norm_before_pred:
+            feat = nn.functional.normalize(feat, dim=-1, p=2)
+        pred = self.predictor(feat)
+        pred = nn.functional.normalize(pred, dim=-1, p=2)
+        return pred
+
+
+class Local_Group_Superivsion_Head(nn.Module):
+    """
+    a class to implement Local Group Superivsion Head which is the same as Instance Superivsion Head
+    --in_dim: input dimension of projection head
+    --hidden_dim: hidden dimension of projection head
+    --out_dim: ouput dimension of projection and prediction heads
+    --pred_hidden_dim: hidden dimension of prediction head
+    --nlayers: layer number of projection head. prediction head has nlayers-1 layer
+    --proj_bn: whether we use batch normalization in projection head
+    --pred_bn: whether we use batch normalization in prediction head
+    --norm_before_pred:  whether we use normalization before prediction head
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        hidden_dim=2048,
+        out_dim=256,
+        pred_hidden_dim=4096,
+        nlayers=3,
+        proj_bn=False,
+        pred_bn=False,
+        norm_before_pred=True,
+    ):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        self.norm_before_pred = norm_before_pred
+
+        self.projector = self._build_mlp(
+            nlayers, in_dim, hidden_dim, out_dim, use_bn=proj_bn
+        )
+
+        self.apply(self._init_weights)
+
+        self.predictor = None
+        if pred_hidden_dim > 0:  # teacher no, student yes
+            self.predictor = self._build_mlp(
+                nlayers - 1, out_dim, pred_hidden_dim, out_dim, use_bn=pred_bn
+            )
+
+    def _init_weights(self, m):
+        """
+        initilize the parameters in network
+        """
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def _build_mlp(self, num_layers, input_dim, hidden_dim, output_dim, use_bn=False):
+        """
+        build a mlp
+        """
+        mlp = []
+        for layer in range(num_layers):
+            dim1 = input_dim if layer == 0 else hidden_dim
+            dim2 = output_dim if layer == num_layers - 1 else hidden_dim
+
+            mlp.append(nn.Linear(dim1, dim2, bias=False))
+
+            if layer < num_layers - 1:
+                if use_bn:
+                    mlp.append(nn.BatchNorm1d(dim2))
+                mlp.append(nn.GELU())
+
+        return nn.Sequential(*mlp)
+
+    def forward(self, x, return_target=False):
+        """
+        forward the input through projection head for teacher and
+        projection/prediction heads for student
+        """
+        feat = self.projector(x)
+
+        if return_target:
+            feat = nn.functional.normalize(feat, dim=-1, p=2)
+            return feat
+        ## return prediction
+        if self.norm_before_pred:
+            feat = nn.functional.normalize(feat, dim=-1, p=2)
+        pred = self.predictor(feat)
+        pred = nn.functional.normalize(pred, dim=-1, p=2)
+        return pred
+
+
+class Group_Superivsion_Head(nn.Module):
+    """
+    a class to implement Local Group Superivsion Head which is the same as Instance Superivsion Head
+    --in_dim: input dimension of projection head
+    --hidden_dim: hidden dimension of projection head
+    --out_dim: ouput dimension of projection and prediction heads
+    --pred_hidden_dim: hidden dimension of prediction head
+    --nlayers: layer number of projection head. prediction head has nlayers-1 layer
+    --proj_bn: whether we use batch normalization in projection head
+    --pred_bn: whether we use batch normalization in prediction head
+    --norm_before_pred:  whether we use normalization before prediction head
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        hidden_dim=2048,
+        bottleneck_dim=256,
+        nlayers=3,
+        use_bn=False,
+        norm_last_layer=True,
+    ):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+
+        self.projector = self._build_mlp(
+            nlayers, in_dim, hidden_dim, bottleneck_dim, use_bn=use_bn
+        )
+        self.apply(self._init_weights)
+
+        self.last_layer = nn.utils.weight_norm(
+            nn.Linear(bottleneck_dim, out_dim, bias=False)
+        )
+        self.last_layer.weight_g.data.fill_(1)
+        if norm_last_layer:
+            self.last_layer.weight_g.requires_grad = False
+
+    def _build_mlp(self, num_layers, in_dim, hidden_dim, output_dim, use_bn=False):
+        """
+        build a mlp
+        """
+        if num_layers == 1:
+            mlp = nn.Linear(in_dim, output_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+            for _ in range(num_layers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if use_bn:
+                    layers.append(nn.BatchNorm1d(hidden_dim))
+                layers.append(nn.GELU())
+            layers.append(nn.Linear(hidden_dim, output_dim))
+            mlp = nn.Sequential(*layers)
+        return mlp
+
+    def _init_weights(self, m):
+        """
+        initilize the parameters in network
+        """
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        """
+        forward the input through the projection and last prediction layer
+        """
+        feat = self.projector(x)
+        feat = nn.functional.normalize(feat, dim=-1, p=2)
+        feat = self.last_layer(feat)
+        return feat
+
+
+class Block_mem(nn.Module):
+    """
+    a class to implement a memory block for local group supervision
+    --dim: feature vector dimenstion in the memory
+    --K: memory size
+    --top_n: number for neighbors in local group supervision
+    """
+
+    def __init__(self, dim, K=2048, top_n=10):
+        super().__init__()
+        self.dim = dim
+        self.K = K
+        self.top_n = top_n
+        # create the queue
+        self.register_buffer("queue_q", torch.randn(K, dim))
+        self.register_buffer("queue_k", torch.randn(K, dim))
+        self.register_buffer("queue_v", torch.randn(K, dim))
+        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self, query, weak_aug_flags):
+        """
+        update memory queue
+        """
+        # import pdb
+        # pdb.set_trace()
+        len_weak = 0
+        query = concat_all_gather(query)
+        if weak_aug_flags is not None:
+            weak_aug_flags = weak_aug_flags.cuda()
+            weak_aug_flags = concat_all_gather(weak_aug_flags)
+            idx_weak = torch.nonzero(weak_aug_flags)
+            len_weak = len(idx_weak)
+            if len_weak > 0:
+                idx_weak = idx_weak.squeeze(-1) 
+                query = query[idx_weak]
+            else:
+                return len_weak
+
+        all_size = query.shape[0]
+        ptr = int(self.queue_ptr)
+        remaining_size = ptr + all_size - self.K
+        if remaining_size <= 0:
+            self.queue_q[ptr : ptr + all_size, :] = query
+            self.queue_k[ptr : ptr + all_size, :] = query
+            self.queue_v[ptr : ptr + all_size, :] = query
+            ptr = ptr + all_size
+            self.queue_ptr[0] = (ptr + all_size) % self.K
+        else:
+            self.queue_q[ptr : self.K, :] = query[0 : self.K - ptr, :]
+            self.queue_k[ptr : self.K, :] = query[0 : self.K - ptr, :]
+            self.queue_v[ptr : self.K, :] = query[0 : self.K - ptr, :]
+
+            self.queue_q[0:remaining_size, :] = query[self.K - ptr :, :]
+            self.queue_k[0:remaining_size, :] = query[self.K - ptr :, :]
+            self.queue_v[0:remaining_size, :] = query[self.K - ptr :, :]
+            self.queue_ptr[0] = remaining_size
+        return len_weak
+
+    @torch.no_grad()
+    def _get_similarity_index(self, x):
+        """
+        compute the index of the top-n neighbors (key-value pair) in memory
+        """
+        x = nn.functional.normalize(x, dim=-1)
+        queue_q = nn.functional.normalize(self.queue_q, dim=-1)
+
+        cosine = x @ queue_q.T
+        _, index = torch.topk(cosine, self.top_n, dim=-1)
+        return index
+
+    @torch.no_grad()
+    def _get_similarity_samples(self, query, index=None):
+        """
+        compute top-n neighbors (key-value pair) in memory
+        """
+        if index is None:
+            index = self._get_similarity_index(query)
+        get_k = self.queue_k[index.view(-1)]
+        get_v = self.queue_v[index.view(-1)]
+        B, tn = index.shape
+        get_k = get_k.view(B, tn, self.dim)
+        get_v = get_v.view(B, tn, self.dim)
+        return get_k, get_v
+
+    def forward(self, query):
+        """
+        forward to find the top-n neighbors (key-value pair) in memory
+        """
+        get_k, get_v = self._get_similarity_samples(query)
+        return get_k, get_v
+
+
+class vit_mem(nn.Module):
+    """
+    a class to implement a memory for local group supervision
+    --dim: feature vector dimenstion in the memory
+    --K: memory size
+    --top_n: number for neighbors in local group supervision
+    """
+
+    def __init__(self, dim, K=2048, top_n=10):
+        super().__init__()
+        self.block = Block_mem(dim, K, top_n)
+
+    def _dequeue_and_enqueue(self, query, weak_aug_flags):
+        """
+        update memory queue
+        """
+        query = query.float()
+        weak_num = self.block._dequeue_and_enqueue(query, weak_aug_flags)
+        return weak_num
+
+    def forward(self, query):
+        """
+        forward to find the top-n neighbors (key-value pair) in memory
+        """
+        query = query.float()
+        get_k, get_v = self.block(query)
+        return get_k, get_v
+
+
+class Mugs_Wrapper(nn.Module):
+    """
+    a class to implement a student or teacher wrapper for mugs
+    --backbone: the backnone of student/teacher, e.g. ViT-small
+    --instance_head: head, including projection/prediction heads, for instance supervision
+    --local_group_head: head, including projection/prediction heads, for local group supervision
+    --group_head: projection head for group supervision
+    """
+
+    def __init__(self, backbone, instance_head, local_group_head, group_head):
+        super(Mugs_Wrapper, self).__init__()
+        backbone.fc, backbone.head = nn.Identity(), nn.Identity()
+        self.backbone = backbone
+        self.instance_head = instance_head
+        self.local_group_head = local_group_head
+        self.group_head = group_head
+
+    def forward(self, x, return_target=False, local_group_memory_inputs=None):
+        """
+        forward input to get instance/local-group/group targets or predictions
+        """
+        # convert to list
+        if not isinstance(x, list):
+            x = [x]
+        idx_crops = torch.cumsum(
+            torch.unique_consecutive(
+                torch.tensor([inp.shape[-1] for inp in x]),
+                return_counts=True,
+            )[1],
+            0,
+        )
+
+        start_idx = 0
+        class_tokens = torch.empty(0).to(x[0].device)
+        mean_patch_tokens = torch.empty(0).to(x[0].device)
+        memory_class_tokens = torch.empty(0).to(x[0].device)
+        for _, end_idx in enumerate(idx_crops):
+            input = torch.cat(x[start_idx:end_idx])
+            token_feat, memory_class_token_feat = self.backbone(
+                input,
+                return_all=True,
+                local_group_memory_inputs=local_group_memory_inputs,
+            )  # [[16, 197, 384], [16, 384]] teacher
+            # [[16, 197, 384], [16, 384]] student  [[48, 37, 384], [48, 384]]
+
+            class_token_feat = token_feat[
+                :, 0
+            ]  # class tokens in ViT, [16, 384] teacher [16, 384] student  [48, 384]
+            class_tokens = torch.cat((class_tokens, class_token_feat))
+
+            start_idx = end_idx
+
+            if self.local_group_head is not None:
+                memory_class_tokens = torch.cat(
+                    (memory_class_tokens, memory_class_token_feat)
+                )
+                if input.shape[-1] == 224:
+                    mean_patch_tokens = torch.cat(
+                        (mean_patch_tokens, token_feat[:, 1:].mean(dim=1))
+                    )
+
+        ## target [16, 256] for teacher,  [64, 256] for student,
+        instance_feat = (
+            self.instance_head(class_tokens, return_target)
+            if self.instance_head is not None
+            else None
+        )
+
+        ## target [16, 256] for teacher,  [64, 256] for student
+        local_group_feat = (
+            self.local_group_head(memory_class_tokens, return_target)
+            if self.local_group_head is not None
+            else None
+        )
+
+        # target [16, 65536] for teacher, [64, 65536] for student
+        group_feat = (
+            self.group_head(class_tokens) if self.group_head is not None else None
+        )
+        return instance_feat, local_group_feat, group_feat, mean_patch_tokens.detach()
+
+
+def get_model(args):
+    """
+    build a student or teacher for mugs, includeing backbone, instance/local-group/group heads,
+    and memory buffer
+    """
+    ## backbone
+    if args.arch in vits.__dict__.keys():
+        student = vits.__dict__[args.arch](
+            patch_size=args.patch_size,
+            num_relation_blocks=1,
+            drop_path_rate=args.drop_path_rate,  # stochastic depth
+        )
+        teacher = vits.__dict__[args.arch](
+            patch_size=args.patch_size, num_relation_blocks=1
+        )
+        embed_dim = student.embed_dim
+    else:
+        assert f"Unknow architecture: {args.arch}"
+
+    ## memory buffer for local-group loss
+    student_mem = vit_mem(
+        embed_dim, K=args.local_group_queue_size, top_n=args.local_group_knn_top_n
+    )
+    teacher_mem = vit_mem(
+        embed_dim, K=args.local_group_queue_size, top_n=args.local_group_knn_top_n
+    )
+
+    ## multi-crop wrapper handles forward with inputs of different resolutions
+    student_instance_head, student_local_group_head, student_group_head = (
+        None,
+        None,
+        None,
+    )
+    teacher_instance_head, teacher_local_group_head, teacher_group_head = (
+        None,
+        None,
+        None,
+    )
+
+    # instance head
+    if args.loss_weights[0] > 0:
+        student_instance_head = Instance_Superivsion_Head(
+            in_dim=embed_dim,
+            hidden_dim=2048,
+            out_dim=args.instance_out_dim,
+            pred_hidden_dim=4096,
+            nlayers=3,
+            proj_bn=args.use_bn_in_head,
+            pred_bn=False,
+            norm_before_pred=args.norm_before_pred,
+        )
+        teacher_instance_head = Instance_Superivsion_Head(
+            in_dim=embed_dim,
+            hidden_dim=2048,
+            out_dim=args.instance_out_dim,
+            pred_hidden_dim=0,
+            nlayers=3,
+            proj_bn=args.use_bn_in_head,
+            pred_bn=False,
+            norm_before_pred=args.norm_before_pred,
+        )
+
+    # local group head
+    if args.loss_weights[1] > 0:
+        student_local_group_head = Local_Group_Superivsion_Head(
+            in_dim=embed_dim,
+            hidden_dim=2048,
+            out_dim=args.local_group_out_dim,
+            pred_hidden_dim=4096,
+            nlayers=3,
+            proj_bn=args.use_bn_in_head,
+            pred_bn=False,
+            norm_before_pred=args.norm_before_pred,
+        )
+        teacher_local_group_head = Local_Group_Superivsion_Head(
+            in_dim=embed_dim,
+            hidden_dim=2048,
+            out_dim=args.local_group_out_dim,
+            pred_hidden_dim=0,
+            nlayers=3,
+            proj_bn=args.use_bn_in_head,
+            pred_bn=False,
+            norm_before_pred=args.norm_before_pred,
+        )
+
+    # group head
+    if args.loss_weights[2] > 0:
+        student_group_head = Group_Superivsion_Head(
+            in_dim=embed_dim,
+            out_dim=args.group_out_dim,
+            hidden_dim=2048,
+            bottleneck_dim=args.group_bottleneck_dim,
+            nlayers=3,
+            use_bn=args.use_bn_in_head,
+            norm_last_layer=args.norm_last_layer,
+        )
+        teacher_group_head = Group_Superivsion_Head(
+            in_dim=embed_dim,
+            out_dim=args.group_out_dim,
+            hidden_dim=2048,
+            bottleneck_dim=args.group_bottleneck_dim,
+            nlayers=3,
+            use_bn=args.use_bn_in_head,
+            norm_last_layer=args.norm_last_layer,
+        )
+
+    # multi-crop wrapper
+    student = Mugs_Wrapper(
+        student, student_instance_head, student_local_group_head, student_group_head
+    )
+
+    teacher = Mugs_Wrapper(
+        teacher, teacher_instance_head, teacher_local_group_head, teacher_group_head
+    )
+
+    return student, teacher, student_mem, teacher_mem
+
+
+# utils
+@torch.no_grad()
+def concat_all_gather(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+
+    tensors_gather = [
+        torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())
+    ]
+    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
+
+    output = torch.cat(tensors_gather, dim=0)
+    return output

src/multicropdataset.py

@@ -0,0 +1,445 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+multi-crop dataset to implement multi-crop augmentation and also dataset
+"""
+import copy
+import random
+
+import torch
+import torchvision.transforms as transforms
+from PIL import Image, ImageFilter, ImageOps
+from src.dataset import ImageFolder
+from src.RandAugment import rand_augment_transform
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.data.random_erasing import RandomErasing
+from timm.data.transforms import _pil_interp
+
+
+class GaussianBlur(object):
+    """
+    Apply Gaussian Blur to the PIL image.
+    """
+
+    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.0):
+        self.prob = p
+        self.radius_min = radius_min
+        self.radius_max = radius_max
+
+    def __call__(self, img):
+        do_it = random.random() <= self.prob
+        if not do_it:
+            return img
+
+        return img.filter(
+            ImageFilter.GaussianBlur(
+                radius=random.uniform(self.radius_min, self.radius_max)
+            )
+        )
+
+
+class Solarization(object):
+    """
+    Apply Solarization to the PIL image.
+    """
+
+    def __init__(self, p):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return ImageOps.solarize(img)
+        else:
+            return img
+
+
+def strong_transforms(
+    img_size=224,
+    scale=(0.08, 1.0),
+    ratio=(0.75, 1.3333333333333333),
+    hflip=0.5,
+    vflip=0.0,
+    color_jitter=0.4,
+    auto_augment="rand-m9-mstd0.5-inc1",
+    interpolation="random",
+    use_prefetcher=True,
+    mean=IMAGENET_DEFAULT_MEAN,  # (0.485, 0.456, 0.406)
+    std=IMAGENET_DEFAULT_STD,  # (0.229, 0.224, 0.225)
+    re_prob=0.25,
+    re_mode="pixel",
+    re_count=1,
+    re_num_splits=0,
+    color_aug=False,
+    strong_ratio=0.45,
+):
+    """
+    for use in a mixing dataset that passes
+     * all data through the first (primary) transform, called the 'clean' data
+     * a portion of the data through the secondary transform
+     * normalizes and converts the branches above with the third, final transform
+    """
+
+    scale = tuple(scale or (0.08, 1.0))  # default imagenet scale range
+    ratio = tuple(ratio or (3.0 / 4.0, 4.0 / 3.0))  # default imagenet ratio range
+
+    primary_tfl = []
+    if hflip > 0.0:
+        primary_tfl += [transforms.RandomHorizontalFlip(p=hflip)]
+    if vflip > 0.0:
+        primary_tfl += [transforms.RandomVerticalFlip(p=vflip)]
+
+    secondary_tfl = []
+    if auto_augment:
+        assert isinstance(auto_augment, str)
+        if isinstance(img_size, tuple):
+            img_size_min = min(img_size)
+        else:
+            img_size_min = img_size
+        aa_params = dict(
+            translate_const=int(img_size_min * strong_ratio),
+            img_mean=tuple([min(255, round(255 * x)) for x in mean]),
+        )
+        if interpolation and interpolation != "random":
+            aa_params["interpolation"] = _pil_interp(interpolation)
+        if auto_augment.startswith("rand"):
+            secondary_tfl += [rand_augment_transform(auto_augment, aa_params)]
+    if color_jitter is not None and color_aug:
+        # color jitter is enabled when not using AA
+        flip_and_color_jitter = [
+            transforms.RandomApply(
+                [
+                    transforms.ColorJitter(
+                        brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1
+                    )
+                ],
+                p=0.8,
+            ),
+            transforms.RandomGrayscale(p=0.2),
+        ]
+        secondary_tfl += flip_and_color_jitter
+
+    if interpolation == "random":
+        interpolation = (Image.BILINEAR, Image.BICUBIC)
+    else:
+        interpolation = _pil_interp(interpolation)
+    final_tfl = [
+        transforms.RandomResizedCrop(
+            size=img_size, scale=scale, ratio=ratio, interpolation=Image.BICUBIC
+        )
+    ]
+    if use_prefetcher:
+        # prefetcher and collate will handle tensor conversion and norm
+        final_tfl += [transforms.ToTensor()]
+    else:
+        final_tfl += [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=torch.tensor(mean), std=torch.tensor(std)),
+        ]
+    if re_prob > 0.0:
+        final_tfl.append(
+            RandomErasing(
+                re_prob,
+                mode=re_mode,
+                max_count=re_count,
+                num_splits=re_num_splits,
+                device="cpu",
+            )
+        )
+    return transforms.Compose(primary_tfl + secondary_tfl + final_tfl)
+
+
+class DataAugmentation(object):
+    """
+    implement multi-crop data augmentation.
+    --global_crops_scale: scale range of the 224-sized cropped image before resizing
+    --local_crops_scale: scale range of the 96-sized cropped image before resizing
+    --local_crops_number: Number of small local views to generate
+    --prob: when we use strong augmentation and weak augmentation, the ratio of images to
+        be cropped with strong augmentation
+    --vanilla_weak_augmentation: whether we use the same augmentation in DINO, namely
+        only using weak augmentation
+    --color_aug: after AutoAugment, whether we further perform color augmentation
+    --local_crop_size: the small crop size
+    --timm_auto_augment_par: the parameters for the AutoAugment used in DeiT
+    --strong_ratio: the ratio of image augmentation for the AutoAugment used in DeiT
+    --re_prob: the re-prob parameter of image augmentation for the AutoAugment used in DeiT
+    --use_prefetcher: whether we use prefetcher which can accerelate the training speed
+    """
+
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        prob=0.5,
+        vanilla_weak_augmentation=False,
+        color_aug=False,
+        local_crop_size=[96],
+        timm_auto_augment_par="rand-m9-mstd0.5-inc1",
+        strong_ratio=0.45,
+        re_prob=0.25,
+        use_prefetcher=False,
+    ):
+
+        ## propability to perform strong augmentation
+        self.prob = prob
+        ## whether we use the commonly used augmentations, e.g. DINO or MoCo-V3
+        self.vanilla_weak_augmentation = vanilla_weak_augmentation
+
+        flip_and_color_jitter = transforms.Compose(
+            [
+                transforms.RandomHorizontalFlip(p=0.5),
+                transforms.RandomApply(
+                    [
+                        transforms.ColorJitter(
+                            brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1
+                        )
+                    ],
+                    p=0.8,
+                ),
+                transforms.RandomGrayscale(p=0.2),
+            ]
+        )
+
+        if use_prefetcher:
+            normalize = transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                ]
+            )
+        else:
+            normalize = transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+                ]
+            )
+
+        ##====== build augmentation of global crops, i.e. 224-sized image crops =========
+        # first global crop, always weak augmentation
+        self.global_transfo1 = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    224, scale=global_crops_scale, interpolation=Image.BICUBIC
+                ),
+                flip_and_color_jitter,
+                GaussianBlur(1.0),
+                normalize,
+            ]
+        )
+
+        # second global crop, always weak augmentation
+        self.global_transfo2 = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    224, scale=global_crops_scale, interpolation=Image.BICUBIC
+                ),
+                flip_and_color_jitter,
+                GaussianBlur(0.1),
+                Solarization(0.2),
+                normalize,
+            ]
+        )
+
+        # strong augmentation, maybe used if we need to perform strong augmentation
+        self.global_transfo3 = strong_transforms(
+            img_size=224,
+            scale=global_crops_scale,
+            ratio=(0.75, 1.3333333333333333),
+            hflip=0.5,
+            vflip=0.0,
+            color_jitter=0.4,
+            auto_augment=timm_auto_augment_par,  # 'rand-m9-mstd0.5-inc1'
+            interpolation="random",
+            use_prefetcher=use_prefetcher,  # True
+            mean=IMAGENET_DEFAULT_MEAN,  # (0.485, 0.456, 0.406)
+            std=IMAGENET_DEFAULT_STD,  # (0.229, 0.224, 0.225)
+            re_prob=re_prob,  # 0.25
+            re_mode="pixel",
+            re_count=1,
+            re_num_splits=0,
+            color_aug=color_aug,
+            strong_ratio=strong_ratio,
+        )
+
+        ##====== build augmentation of local crops, i.e. 96-sized image crops =========
+        self.local_crops_number = (
+            local_crops_number  # transformation for the local small crops
+        )
+        assert local_crop_size[0] == 96
+        # weak augmentation, maybe used if we need to perform weak augmentation
+        self.local_transfo = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    local_crop_size[0],
+                    scale=local_crops_scale,
+                    interpolation=Image.BICUBIC,
+                ),
+                flip_and_color_jitter,
+                GaussianBlur(p=0.5),
+                normalize,
+            ]
+        )
+        # strong augmentation, maybe used if we need to perform strong augmentation
+        self.local_transfo2 = strong_transforms(
+            img_size=local_crop_size[0],  # (224, 224)
+            scale=local_crops_scale,  # (0.08, 1.0)
+            ratio=(0.75, 1.3333333333333333),  # (0.75, 1.3333333333333333)
+            hflip=0.5,  # 0.5
+            vflip=0.0,  # 0.0
+            color_jitter=0.4,  # 0.4
+            auto_augment=timm_auto_augment_par,  # 'rand-m9-mstd0.5-inc1'
+            interpolation="random",  # 'random'
+            use_prefetcher=use_prefetcher,  # True
+            mean=IMAGENET_DEFAULT_MEAN,  # (0.485, 0.456, 0.406)
+            std=IMAGENET_DEFAULT_STD,  # (0.229, 0.224, 0.225)
+            re_prob=re_prob,  # 0.25
+            re_mode="pixel",  # 'pixel'
+            re_count=1,  # 1
+            re_num_splits=0,  # 0
+            color_aug=color_aug,
+            strong_ratio=strong_ratio,
+        )
+
+    def __call__(self, image):
+        """
+        implement multi-crop data augmentation. Generate two 224-sized +
+        "local_crops_number" 96-sized images
+        """
+        crops = []
+        ##====== images to be fed into teacher, two 224-sized =========
+        img1 = self.global_transfo1(image)
+        img2 = self.global_transfo2(image)
+        crops.append(img1)
+        crops.append(img2)
+
+        ##====== images to be fed into student, two 224-sized + "local_crops_number" 96-sized =========
+        # first to generate two 224-sized
+        # this weak_flag indicates whether the current image is weakly augmented.
+        # For local group supervision, we only use weakly augmented images of size 224 to
+        # update the memory for local-group aggregation.
+        weak_flag = False
+
+        if self.vanilla_weak_augmentation is True:
+            ## directly copy the images of weak augmentation
+            crops.append(copy.deepcopy(img1))
+            crops.append(copy.deepcopy(img2))
+            weak_flag = True
+        elif self.prob < 1.0 and random.random() > self.prob:
+            ## whether perform strong augmentation
+            crops.append(self.global_transfo3(image))
+            crops.append(self.global_transfo3(image))
+        else:
+            ## perform weak augmentation
+            crops.append(self.global_transfo1(image))
+            crops.append(self.global_transfo2(image))
+            weak_flag = True
+
+        # then to generate "local_crops_number" 96-sized
+        for _ in range(self.local_crops_number):
+            if self.prob < 1.0 and random.random() > self.prob:
+                ## whether perform strong augmentation
+                crops.append(self.local_transfo2(image))
+            else:
+                ## perform weak augmentation
+                crops.append(self.local_transfo(image))
+
+        return crops, weak_flag
+
+
+def get_dataset(args):
+    """
+    build a multi-crop data augmentation and a dataset/dataloader
+    """
+    ## preparing augmentations, including weak and strong augmentations
+    transform = DataAugmentation(
+        global_crops_scale=args.global_crops_scale,
+        local_crops_scale=args.local_crops_scale,
+        local_crops_number=args.local_crops_number,
+        vanilla_weak_augmentation=args.vanilla_weak_augmentation,
+        prob=args.prob,
+        color_aug=args.color_aug,
+        local_crop_size=args.size_crops,
+        timm_auto_augment_par=args.timm_auto_augment_par,
+        strong_ratio=args.strong_ratio,
+        re_prob=args.re_prob,
+        use_prefetcher=args.use_prefetcher,
+    )
+
+    ## For debug mode, we only load the first two classes to reduce data reading time.
+    ## otherwise, we load all training data for pretraining.
+    class_num = 2 if args.debug else 1000
+    dataset = ImageFolder(args.data_path, transform=transform, class_num=class_num)
+
+    sampler = torch.utils.data.DistributedSampler(dataset, shuffle=True)
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        sampler=sampler,
+        batch_size=args.batch_size_per_gpu,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+    )
+    return data_loader
+
+
+class data_prefetcher:
+    """
+    implement data prefetcher. we perform some augmentation on GPUs intead of CPUs
+    --loader: a data loader
+    --fp16: whether we use fp16, if yes, we need to tranform the data to be fp16
+    """
+
+    def __init__(self, loader, fp16=True):
+        self.loader = iter(loader)
+        self.fp16 = fp16
+        self.stream = torch.cuda.Stream()
+        self.mean = torch.tensor([0.485, 0.456, 0.406]).cuda().view(1, 3, 1, 1)
+        self.std = torch.tensor([0.229, 0.224, 0.225]).cuda().view(1, 3, 1, 1)
+        if fp16:
+            self.mean = self.mean.half()
+            self.std = self.std.half()
+
+        self.preload()
+
+    def preload(self):
+        """
+        preload the next minibatch of data
+        """
+        try:
+            self.multi_crops, self.weak_flag = next(self.loader)
+        except StopIteration:
+            self.multi_crops, self.weak_flag = None, None
+            return
+
+        with torch.cuda.stream(self.stream):
+            for i in range(len(self.multi_crops)):
+                self.multi_crops[i] = self.multi_crops[i].cuda(non_blocking=True)
+                if self.fp16:
+                    self.multi_crops[i] = (
+                        self.multi_crops[i].half().sub_(self.mean).div_(self.std)
+                    )
+                else:
+                    self.multi_crops[i] = (
+                        self.multi_crops[i].float().sub_(self.mean).div_(self.std)
+                    )
+
+    def next(self):
+        """
+        load the next minibatch of data
+        """
+        torch.cuda.current_stream().wait_stream(self.stream)
+        multi_crops, weak_flags = self.multi_crops, self.weak_flag
+        self.preload()
+        return multi_crops, weak_flags

src/optimizer.py

@@ -0,0 +1,210 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+implment some functions for optimizers
+"""
+import numpy as np
+import torch
+
+import utils
+
+
+def clip_gradients(model, clip):
+    """
+    clip gradient if gradient norm > clip
+    """
+    norms = []
+    for name, p in model.named_parameters():
+        if p.grad is not None:
+            param_norm = p.grad.data.norm(2)
+            norms.append(param_norm.item())
+            clip_coef = clip / (param_norm + 1e-6)
+            if clip_coef < 1:
+                p.grad.data.mul_(clip_coef)
+    return norms
+
+
+def cancel_gradients_last_layer(epoch, model, freeze_last_layer):
+    """
+    cancle gradient if epoch > freeze_last_layer
+    """
+    if epoch >= freeze_last_layer:
+        return
+    for n, p in model.named_parameters():
+        if "last_layer" in n:
+            p.grad = None
+
+
+def cosine_scheduler(
+    base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, start_warmup_value=0
+):
+    """
+    start_warmup_value to base_value in the first warmup_epochs epochs;
+    then cosine scheduling base_value to final_value in the remaining epochs-warmup_epochs
+    """
+    warmup_schedule = np.array([])
+    warmup_iters = warmup_epochs * niter_per_ep
+    if warmup_epochs > 0:
+        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+    iters = np.arange(epochs * niter_per_ep - warmup_iters)
+    schedule = final_value + 0.5 * (base_value - final_value) * (
+        1 + np.cos(np.pi * iters / len(iters))
+    )
+
+    schedule = np.concatenate((warmup_schedule, schedule))
+    assert len(schedule) == epochs * niter_per_ep
+    return schedule
+
+
+def get_params_groups(model):
+    """
+    divide the parameters into several groups, see below
+    """
+    regularized = []
+    not_regularized = []
+    patch_embed = []
+    patch_embed_not_regularized = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        # we do not regularize biases nor Norm parameters
+        if name.endswith(".bias") or len(param.shape) == 1:
+            if "patch_embed" in name:
+                patch_embed_not_regularized.append(param)
+            else:
+                not_regularized.append(param)
+        elif "patch_embed" in name:
+            patch_embed.append(param)
+        else:
+            regularized.append(param)
+    return [
+        {"name": "normal_params", "params": regularized},
+        {"name": "patch_embed", "params": patch_embed},
+        {
+            "name": "no_wd",
+            "params": not_regularized,
+            "apply_wd": False,
+            "weight_decay": 0.0,
+        },
+        {
+            "name": "patch_embed_no_wd",
+            "params": patch_embed_not_regularized,
+            "apply_wd": False,
+            "weight_decay": 0.0,
+        },
+    ]
+
+
+class LARS(torch.optim.Optimizer):
+    """
+    Almost copy-paste from https://github.com/facebookresearch/barlowtwins/blob/main/main.py
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=0,
+        weight_decay=0,
+        momentum=0.9,
+        eta=0.001,
+        weight_decay_filter=None,
+        lars_adaptation_filter=None,
+    ):
+        defaults = dict(
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=momentum,
+            eta=eta,
+            weight_decay_filter=weight_decay_filter,
+            lars_adaptation_filter=lars_adaptation_filter,
+        )
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for g in self.param_groups:
+            for p in g["params"]:
+                dp = p.grad
+
+                if dp is None:
+                    continue
+
+                if p.ndim != 1:
+                    dp = dp.add(p, alpha=g["weight_decay"])
+
+                if p.ndim != 1:
+                    param_norm = torch.norm(p)
+                    update_norm = torch.norm(dp)
+                    one = torch.ones_like(param_norm)
+                    q = torch.where(
+                        param_norm > 0.0,
+                        torch.where(
+                            update_norm > 0, (g["eta"] * param_norm / update_norm), one
+                        ),
+                        one,
+                    )
+                    dp = dp.mul(q)
+
+                param_state = self.state[p]
+                if "mu" not in param_state:
+                    param_state["mu"] = torch.zeros_like(p)
+                mu = param_state["mu"]
+                mu.mul_(g["momentum"]).add_(dp)
+
+                p.add_(mu, alpha=-g["lr"])
+
+
+def get_optimizer(student, len_dataloader, args):
+    """
+    build an optimizer for training
+    """
+    # ============ preparing optimizer ... ============
+    params_groups = get_params_groups(student)
+    if args.optimizer == "adamw":
+        optimizer = torch.optim.AdamW(params_groups)  # to use with ViTs
+    elif args.optimizer == "sgd":
+        optimizer = torch.optim.SGD(
+            params_groups, lr=0, momentum=0.9
+        )  # lr is set by scheduler
+    elif args.optimizer == "lars":
+        optimizer = LARS(params_groups)  # to use with convnet and large batches
+    # for mixed precision training
+    fp16_scaler = None
+    if args.use_fp16:
+        fp16_scaler = torch.cuda.amp.GradScaler()
+
+    # ============ init schedulers ... ============
+    lr_schedule = cosine_scheduler(
+        args.lr
+        * (args.batch_size_per_gpu * utils.get_world_size())
+        / 256.0,  # linear scaling rule
+        args.min_lr,
+        args.epochs,
+        len_dataloader,
+        warmup_epochs=args.warmup_epochs,
+    )
+    wd_schedule = cosine_scheduler(
+        args.weight_decay,
+        args.weight_decay_end,
+        args.epochs,
+        len_dataloader,  # len(data_loader),
+    )
+    # momentum parameter is increased to 1. during training with a cosine schedule
+    momentum_schedule = cosine_scheduler(
+        args.momentum_teacher, 1, args.epochs, len_dataloader
+    )
+    print("Loss, optimizer and schedulers ready.")
+
+    return optimizer, fp16_scaler, lr_schedule, wd_schedule, momentum_schedule

src/vision_transformer.py

@@ -0,0 +1,491 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+ViT backbones, including ViT-small, ViT-base, ViT-large
+Mostly copy-paste from timm library.
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+"""
+import math
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        lower = norm_cdf((a - mean) / std)
+        upper = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * lower - 1, 2 * upper - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.0))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    # type: (torch.tensor, float, float, float, float) -> torch.tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    """
+    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """
+    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    """
+    MLP module in ViT
+    """
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    """
+    Attention module in ViT
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        reshaped_qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = reshaped_qkv[0], reshaped_qkv[1], reshaped_qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+
+class Block(nn.Module):
+    """
+    ViT block, including Attention, MLP, etc.
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, x):
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer"""
+
+    def __init__(
+        self,
+        img_size=[224, 224],
+        patch_size=16,
+        in_chans=3,
+        num_classes=0,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        num_relation_blocks=0,
+        **kwargs
+    ):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+        self.patch_size = patch_size
+        self.num_classes = num_classes
+        self.depth = depth
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0],
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        trunc_normal_(self.pos_embed, std=0.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        # Classifier head
+        self.head = (
+            nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        )
+
+        self.num_relation_blocks = num_relation_blocks
+        if num_relation_blocks > 0:
+            self.relation_blocks = nn.ModuleList(
+                [
+                    Block(
+                        dim=embed_dim,
+                        num_heads=num_heads,
+                        mlp_ratio=mlp_ratio,
+                        qkv_bias=qkv_bias,
+                        qk_scale=qk_scale,
+                        drop=drop_rate,
+                        attn_drop=attn_drop_rate,
+                        drop_path=dpr[i],
+                        norm_layer=norm_layer,
+                    )
+                    for i in range(int(num_relation_blocks))
+                ]
+            )
+
+        trunc_normal_(self.cls_token, std=0.02)
+        self.apply(self._init_weights)
+
+    def add_pos_emb_for_cls_token(self):
+        pe_cls_token = torch.zeros([1, 1, self.embed_dim], dtype=torch.float32)
+        self.pos_embed = nn.Parameter(torch.cat([pe_cls_token, self.pos_embed], dim=1))
+        self.pos_embed.requires_grad = False
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(
+                1, int(math.sqrt(N)), int(math.sqrt(N)), dim
+            ).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode="bicubic",
+        )
+        assert (
+            int(w0) == patch_pos_embed.shape[-2]
+            and int(h0) == patch_pos_embed.shape[-1]
+        )
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)  # patch linear embedding
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, w, h)
+        return self.pos_drop(x)
+
+    def forward(self, x, return_all=False, local_group_memory_inputs=None, **kwargs):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+
+        if self.num_relation_blocks > 0:
+            mem = local_group_memory_inputs.get("mem")
+            if mem is not None:
+                m, _ = mem(x.mean(1))
+                rx = torch.cat((x.mean(1).unsqueeze(1), m), dim=1)
+            else:
+                rx = x
+            for i, blk in enumerate(self.relation_blocks):
+                rx = blk(rx)
+            relation_out = self.norm(rx[:, 0])
+
+        x = self.norm(x)
+        if self.num_classes > 0:
+            return self.head(x[:, 0])
+
+        if return_all:
+            return x, relation_out
+        else:
+            return x[:, 0], relation_out
+
+    def forward_knn(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x[:, 0]
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+    def get_num_layers(self):
+        return self.depth
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+
+def vit_tiny(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model
+
+
+def vit_small(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model
+
+
+def vit_large(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model

utils.py

@@ -0,0 +1,583 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Misc functions.
+
+Mostly copy-paste from torchvision references or other public repos like DETR and DINO:
+https://github.com/facebookresearch/detr/blob/master/util/misc.py
+https://github.com/facebookresearch/dino/blob/main/utils.py
+"""
+import datetime
+import logging
+import os
+import subprocess
+import sys
+import time
+from collections import defaultdict, deque
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch import nn
+
+
+def get_logger(file_path_name):
+    """
+    build a logger which both write on the desk and also on the terminal
+    """
+    logger = logging.getLogger()
+    logger.setLevel("INFO")
+    BASIC_FORMAT = "%(levelname)s:%(message)s"
+    DATE_FORMAT = ""
+    formatter = logging.Formatter(BASIC_FORMAT, DATE_FORMAT)
+    chlr = logging.StreamHandler()
+    chlr.setFormatter(formatter)
+    chlr.setLevel("INFO")
+    fhlr = logging.FileHandler(file_path_name)
+    fhlr.setFormatter(formatter)
+    logger.addHandler(chlr)
+    logger.addHandler(fhlr)
+
+    return logger
+
+
+def restart_from_checkpoint(ckp_path, run_variables=None, **kwargs):
+    """
+    Re-start from checkpoint
+    """
+    if not os.path.isfile(ckp_path):
+        return
+    print("Found checkpoint at {}".format(ckp_path))
+
+    # open checkpoint file
+    checkpoint = torch.load(ckp_path, map_location="cpu")
+    # key is what to look for in the checkpoint file
+    # value is the object to load
+    # example: {'state_dict': model}
+    for key, value in kwargs.items():
+        if key in checkpoint and value is not None:
+            try:
+                msg = value.load_state_dict(checkpoint[key], strict=False)
+                print(
+                    "=> loaded '{}' from checkpoint '{}' with msg {}".format(
+                        key, ckp_path, msg
+                    )
+                )
+            except TypeError:
+                try:
+                    msg = value.load_state_dict(checkpoint[key])
+                    print("=> loaded '{}' from checkpoint: '{}'".format(key, ckp_path))
+                except ValueError:
+                    print(
+                        "=> failed to load '{}' from checkpoint: '{}'".format(
+                            key, ckp_path
+                        )
+                    )
+        else:
+            print("=> key '{}' not found in checkpoint: '{}'".format(key, ckp_path))
+
+    # reload variable important for the run
+    if run_variables is not None:
+        for var_name in run_variables:
+            if var_name in checkpoint:
+                run_variables[var_name] = checkpoint[var_name]
+
+
+def bool_flag(s):
+    """
+    Parse boolean arguments from the command line.
+    """
+    FALSY_STRINGS = {"off", "false", "0"}
+    TRUTHY_STRINGS = {"on", "true", "1"}
+    if s.lower() in FALSY_STRINGS:
+        return False
+    elif s.lower() in TRUTHY_STRINGS:
+        return True
+    else:
+        raise argparse.ArgumentTypeError("invalid value for a boolean flag")
+
+
+def fix_random_seeds(seed=31):
+    """
+    Fix random seeds.
+    """
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+
+
+def has_batchnorms(model):
+    """
+    judge whether a model has batch normalization
+    """
+    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
+    for name, module in model.named_modules():
+        if isinstance(module, bn_types):
+            return True
+    return False
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.6f} ({global_avg:.6f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value,
+        )
+
+
+class MetricLogger(object):
+    """
+    build a Metric Logger
+    """
+
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(
+            "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
+        )
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append("{}: {}".format(name, str(meter)))
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.6f}")
+        data_time = SmoothedValue(fmt="{avg:.6f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
+        else:
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                ]
+            )
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(
+            "{} Total time: {} ({:.6f} s / it)".format(
+                header, total_time_str, total_time / len(iterable)
+            )
+        )
+
+
+def get_sha():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+
+    def _run(command):
+        return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
+
+    sha = "N/A"
+    diff = "clean"
+    branch = "N/A"
+    try:
+        sha = _run(["git", "rev-parse", "HEAD"])
+        subprocess.check_output(["git", "diff"], cwd=cwd)
+        diff = _run(["git", "diff-index", "HEAD"])
+        diff = "has uncommited changes" if diff else "clean"
+        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+    except Exception:
+        pass
+    message = f"sha: {sha}, status: {diff}, branch: {branch}"
+    return message
+
+
+def is_dist_avail_and_initialized():
+    """
+    judge whether distributed training is available and well-initialized
+    """
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    """
+    get the world size
+    """
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    """
+    get the rank
+    """
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    """
+    judge whether the current node is the master node
+    """
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    """
+    save checkpoint on the master node
+    """
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def init_distributed_ddpjob(args=None):
+    """
+    initialize the ddp job
+    """
+    if dist.is_available() and dist.is_initialized():
+        return dist.get_world_size(), dist.get_rank()
+
+    try:
+        os.environ["MASTER_PORT"] = "40101"
+        torch.distributed.init_process_group(backend="nccl")
+    except Exception:
+        world_size, rank = 1, 0
+        print("distributed training not available")
+
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    args.gpu = args.rank
+    args.world_size, args.rank = world_size, rank
+    return world_size, rank
+
+
+def init_distributed_mode(args):
+    """
+    initialize the normal job
+    """
+    # launched with torch.distributed.launch
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ.get("LOCAL_RANK", 0))
+        print(
+            "args.rank",
+            args.rank,
+            "args.world_size",
+            args.world_size,
+            "args.gpu",
+            args.gpu,
+        )
+        print("get_rank()", get_rank())
+    # launched with submitit on a slurm cluster
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    # launched naively with `python main_dino.py`
+    # we manually add MASTER_ADDR and MASTER_PORT to env variables
+    elif torch.cuda.is_available():
+        print("Will run the code on one GPU.")
+        args.rank, args.gpu, args.world_size = 0, 0, 1
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = "2950"
+    else:
+        print("Does not support training without GPU.")
+        sys.exit(1)
+
+    os.environ["MASTER_PORT"] = "6542"
+
+    dist.init_process_group(
+        backend="nccl",
+        init_method=args.dist_url,
+        world_size=args.world_size,
+        rank=args.rank,
+    )
+
+    torch.cuda.set_device(args.gpu)
+    print(
+        "| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True
+    )
+    dist.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+def accuracy(output, target, topk=(1,)):
+    """
+    Computes the accuracy over the k top predictions for the specified values of k
+    """
+    maxk = max(topk)
+    batch_size = target.size(0)
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.reshape(1, -1).expand_as(pred))
+    return [correct[:k].reshape(-1).float().sum(0) * 100.0 / batch_size for k in topk]
+
+
+def multi_scale(samples, model):
+    """
+    build a multi-scale features
+    """
+    v = None
+    for s in [1, 1 / 2 ** (1 / 2), 1 / 2]:  # we use 3 different scales
+        if s == 1:
+            inp = samples.clone()
+        else:
+            inp = nn.functional.interpolate(
+                samples, scale_factor=s, mode="bilinear", align_corners=False
+            )
+        feats = model.forward_knn(inp).clone()
+        if v is None:
+            v = feats
+        else:
+            v += feats
+    v /= 3
+    v /= v.norm()
+    return v
+
+
+class AllGather(torch.autograd.Function):
+    """
+    gather the variable on different nodes toghther
+    """
+
+    @staticmethod
+    def forward(ctx, x):
+        if (
+            dist.is_available()
+            and dist.is_initialized()
+            and (dist.get_world_size() > 1)
+        ):
+            outputs = [torch.zeros_like(x) for _ in range(dist.get_world_size())]
+            dist.all_gather(outputs, x)
+            return torch.cat(outputs, 0)
+        return x
+
+    @staticmethod
+    def backward(ctx, grads):
+        if (
+            dist.is_available()
+            and dist.is_initialized()
+            and (dist.get_world_size() > 1)
+        ):
+            s = (grads.shape[0] // dist.get_world_size()) * dist.get_rank()
+            e = (grads.shape[0] // dist.get_world_size()) * (dist.get_rank() + 1)
+            grads = grads.contiguous()
+            dist.all_reduce(grads)
+            return grads[s:e]
+        return grads
+
+
+class AllReduce(torch.autograd.Function):
+    """
+    reduce the variable on different nodes toghther
+    """
+
+    @staticmethod
+    def forward(ctx, x):
+        if (
+            dist.is_available()
+            and dist.is_initialized()
+            and (dist.get_world_size() > 1)
+        ):
+            x = x.contiguous() / dist.get_world_size()
+            dist.all_reduce(x)
+        return x
+
+    @staticmethod
+    def backward(ctx, grads):
+        return grads
+
+
+def load_pretrained_weights(
+    model, pretrained_weights, checkpoint_key, model_name, patch_size
+):
+    if os.path.isfile(pretrained_weights):
+        state_dict = torch.load(pretrained_weights, map_location="cpu")
+        if checkpoint_key is not None and checkpoint_key in state_dict:
+            print(f"Take key {checkpoint_key} in provided checkpoint dict")
+            state_dict = state_dict[checkpoint_key]
+        # remove `module.` prefix
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        # remove `backbone.` prefix induced by multicrop wrapper
+        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+        # remove `encoder.` prefix induced by MAE
+        state_dict = {k.replace("encoder.", ""): v for k, v in state_dict.items()}
+        msg = model.load_state_dict(state_dict, strict=False)
+        print(
+            "Pretrained weights found at {} and loaded with msg: {}".format(
+                pretrained_weights, msg
+            )
+        )
+    else:
+        print(
+            "There is no reference weights available for this model => We use random weights."
+        )
+
+
+@torch.no_grad()
+def concat_all_gather(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+    tensors_gather = [
+        torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())
+    ]
+    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
+
+    output = torch.cat(tensors_gather, dim=0)
+    return output