Added support for H5 data, improved scripts and data handling

.github/workflows/dali_tests.yml

@@ -34,7 +34,8 @@ jobs:
       - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e .[umap] codecov mypy pytest-cov black
+          pip install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist codecov
+          pip install mypy pytest-cov black
 
       - name: Cache datasets
         uses: actions/cache@v2
@@ -61,4 +62,4 @@ jobs:
           file: coverage.xml
           flags: dali
           name: DALI-coverage
-          fail_ci_if_error: false
+          fail_ci_if_error: false

.github/workflows/tests.yml

@@ -34,7 +34,7 @@ jobs:
       - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e .[umap] codecov mypy pytest-cov black
+          pip install -e .[umap,h5] codecov mypy pytest-cov black
 
       - name: Cache datasets
         uses: actions/cache@v2

README.md

@@ -15,8 +15,9 @@ The library is self-contained, but it is possible to use the models outside of s
 ---
 
 ## News
+* **[Jul 13 2022]**: :sparkling_heart: Added support for [H5](https://docs.h5py.org/en/stable/index.html) data, improved scripts and data handling.
 * **[Jun 26 2022]**: :fire: Added [MoCo V3](https://arxiv.org/abs/2104.02057).
-* **[Jun 10 2022]**: :bomb: Improved LARS and fixed some issues to support [Horovod](https://horovod.readthedocs.io/en/stable/pytorch.html).
+* **[Jun 10 2022]**: :bomb: Improved LARS.
 * **[Jun 09 2022]**: :lollipop: Added support for [WideResnet](https://arxiv.org/abs/1605.07146), multicrop for SwAV and equalization data augmentation.
 * **[May 02 2022]**: :diamond_shape_with_a_dot_inside: Wrapped Dali with a DataModule, added auto resume for linear eval and Wandb run resume.
 * **[Apr 12 2022]**: :rainbow: Improved design of models and added support to train with a fraction of data.
@@ -65,7 +66,7 @@ The library is self-contained, but it is possible to use the models outside of s
 
 ## Extra flavor
 
-### Multiple backbones
+### Backbones
 * [ResNet](https://arxiv.org/abs/1512.03385)
 * [WideResNet](https://arxiv.org/abs/1605.07146)
 * [ViT](https://arxiv.org/abs/2010.11929)
@@ -77,22 +78,23 @@ The library is self-contained, but it is possible to use the models outside of s
 * Increased data processing speed by up to 100% using [Nvidia Dali](https://github.com/NVIDIA/DALI).
 * Flexible augmentations.
 
-### Evaluation and logging
+### Evaluation
 * Online linear evaluation via stop-gradient for easier debugging and prototyping (optionally available for the momentum backbone as well).
+* Standard offline linear evaluation.
 * Online and offline K-NN evaluation.
-* Normal offline linear evaluation.
-* All the perks of PyTorch Lightning (mixed precision, gradient accumulation, clipping, automatic logging and much more).
-* Easy-to-extend modular code structure.
-* Custom model logging with a simpler file organization.
 * Automatic feature space visualization with UMAP.
-* Offline UMAP.
-* Common metrics.
 
 ### Training tricks
+* All the perks of PyTorch Lightning (mixed precision, gradient accumulation, clipping, and much more).
+* Channel last conversion
 * Multi-cropping dataloading following [SwAV](https://arxiv.org/abs/2006.09882):
     * **Note**: currently, only SimCLR, BYOL and SwAV support this.
-* Exclude batchnorm and biases from LARS.
-* No LR scheduler for the projection head in SimSiam.
+* Exclude batchnorm and biases from weight decay and LARS.
+* No LR scheduler for the projection head (as in SimSiam).
+
+### Logging
+* Metric logging on the cloud with [WandB](https://wandb.ai/site)
+* Custom model checkpointing with a simple file organization.
 
 ---
 ## Requirements
@@ -122,10 +124,10 @@ First clone the repo.
 
 Then, to install solo-learn with [Dali](https://github.com/NVIDIA/DALI) and/or UMAP support, use:
 ```
-pip3 install .[dali,umap] --extra-index-url https://developer.download.nvidia.com/compute/redist
+pip3 install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist
 ```
 
-If no Dali/UMAP support is needed, the repository can be installed as:
+If no Dali/UMAP/H5 support is needed, the repository can be installed as:
 ```
 pip3 install .
 ```

main_knn.py

@@ -1,4 +1,4 @@
-# Copyright 2021 solo-learn development team.
+# Copyright 2022 solo-learn development team.
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of
 # this software and associated documentation files (the "Software"), to deal in

main_linear.py

@@ -1,4 +1,4 @@
-# Copyright 2021 solo-learn development team.
+# Copyright 2022 solo-learn development team.
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of
 # this software and associated documentation files (the "Software"), to deal in
@@ -57,7 +57,7 @@ def main():
     if "swin" in args.backbone and cifar:
         kwargs["window_size"] = 4
 
-    backbone = backbone_model(**kwargs)
+    backbone = backbone_model(method=None, **kwargs)
     if args.backbone.startswith("resnet"):
         # remove fc layer
         backbone.fc = nn.Identity()
@@ -90,25 +90,28 @@ def main():
     model = LinearModel(backbone, **args.__dict__)
     make_contiguous(model)
 
+    if args.data_format == "dali":
+        val_data_format = "image_folder"
+    else:
+        val_data_format = args.data_format
     train_loader, val_loader = prepare_data(
         args.dataset,
-        data_dir=args.data_dir,
-        train_dir=args.train_dir,
-        val_dir=args.val_dir,
+        train_data_path=args.train_data_path,
+        val_data_path=args.val_data_path,
+        data_format=val_data_format,
         batch_size=args.batch_size,
         num_workers=args.num_workers,
-        data_fraction=args.data_fraction,
     )
-    if args.dali:
+
+    if args.data_format == "dali":
         assert (
             _dali_avaliable
         ), "Dali is not currently avaiable, please install it first with [dali]."
 
         dali_datamodule = ClassificationDALIDataModule(
             dataset=args.dataset,
-            data_dir=args.data_dir,
-            train_dir=args.train_dir,
-            val_dir=args.val_dir,
+            train_data_path=args.train_data_path,
+            val_data_path=args.val_data_path,
             num_workers=args.num_workers,
             batch_size=args.batch_size,
             data_fraction=args.data_fraction,
@@ -192,7 +195,7 @@ def prefetch_batches(self) -> int:
     except:
         pass
 
-    if args.dali:
+    if args.data_format == "dali":
         trainer.fit(model, ckpt_path=ckpt_path, datamodule=dali_datamodule)
     else:
         trainer.fit(model, train_loader, val_loader, ckpt_path=ckpt_path)

main_pretrain.py

@@ -1,4 +1,4 @@
-# Copyright 2021 solo-learn development team.
+# Copyright 2022 solo-learn development team.
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of
 # this software and associated documentation files (the "Software"), to deal in
@@ -69,28 +69,32 @@ def main():
     make_contiguous(model)
 
     # validation dataloader for when it is available
-    if args.dataset == "custom" and (args.no_labels or args.val_dir is None):
+    if args.dataset == "custom" and (args.no_labels or args.val_data_path is None):
         val_loader = None
-    elif args.dataset in ["imagenet100", "imagenet"] and args.val_dir is None:
+    elif args.dataset in ["imagenet100", "imagenet"] and (args.val_data_path is None):
         val_loader = None
     else:
+        if args.data_format == "dali":
+            val_data_format = "image_folder"
+        else:
+            val_data_format = args.data_format
+
         _, val_loader = prepare_data_classification(
             args.dataset,
-            data_dir=args.data_dir,
-            train_dir=args.train_dir,
-            val_dir=args.val_dir,
+            train_data_path=args.train_data_path,
+            val_data_path=args.val_data_path,
+            data_format=val_data_format,
             batch_size=args.batch_size,
             num_workers=args.num_workers,
         )
 
     # pretrain dataloader
-    if args.dali:
+    if args.data_format == "dali":
         assert _dali_avaliable, "Dali is not avaiable, please install it first with [dali]."
 
         dali_datamodule = PretrainDALIDataModule(
             dataset=args.dataset,
-            data_dir=args.data_dir,
-            train_dir=args.train_dir,
+            train_data_path=args.train_data_path,
             unique_augs=args.unique_augs,
             transform_kwargs=args.transform_kwargs,
             num_crops_per_aug=args.num_crops_per_aug,
@@ -120,8 +124,8 @@ def main():
         train_dataset = prepare_datasets(
             args.dataset,
             transform,
-            data_dir=args.data_dir,
-            train_dir=args.train_dir,
+            train_data_path=args.train_data_path,
+            data_format=args.data_format,
             no_labels=args.no_labels,
             data_fraction=args.data_fraction,
         )
@@ -214,7 +218,7 @@ def prefetch_batches(self) -> int:
     except:
         pass
 
-    if args.dali:
+    if args.data_format == "dali":
         trainer.fit(model, ckpt_path=ckpt_path, datamodule=dali_datamodule)
     else:
         trainer.fit(model, train_loader, val_loader, ckpt_path=ckpt_path)

main_umap.py

@@ -1,4 +1,4 @@
-# Copyright 2021 solo-learn development team.
+# Copyright 2022 solo-learn development team.
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy of
 # this software and associated documentation files (the "Software"), to deal in

requirements.txt

@@ -8,4 +8,4 @@ tqdm
 wandb
 scipy
 timm
-scikit-learn
+scikit-learn

bash_files/knn/imagenet-100/knn.sh → scripts/knn/imagenet-100/knn.sh

@@ -1,8 +1,7 @@
 python3 main_knn.py \
     --dataset imagenet100 \
-    --data_dir /datasets \
-    --train_dir imagenet-100/train \
-    --val_dir imagenet-100/val \
+    --train_data_path /datasets/imagenet-100/train \
+    --val_data_path /datasets/imagenet-100/val \
     --batch_size 16 \
     --num_workers 10 \
     --pretrained_checkpoint_dir PATH \

bash_files/linear/imagenet-100/barlow_linear.sh → scripts/linear/imagenet-100/barlow_linear.sh

@@ -1,9 +1,8 @@
 python3 main_linear.py \
     --dataset imagenet100 \
     --backbone resnet18 \
-    --data_dir /datasets \
-    --train_dir imagenet-100/train \
-    --val_dir imagenet-100/val \
+    --train_data_path /datasets/imagenet-100/train \
+    --val_data_path /datasets/imagenet-100/val \
     --max_epochs 100 \
     --devices 0 \
     --accelerator gpu \
@@ -15,11 +14,11 @@ python3 main_linear.py \
     --weight_decay 0 \
     --batch_size 256 \
     --num_workers 4 \
-    --dali \
+    --data_format dali \
     --name barlow-imagenet100-linear-eval \
     --pretrained_feature_extractor PATH \
     --project solo-learn \
     --entity unitn-mhug \
     --wandb \
     --save_checkpoint \
-    --auto_resume
+    --auto_resume

bash_files/linear/imagenet-100/byol_linear.sh → scripts/linear/imagenet-100/byol_linear.sh

@@ -1,9 +1,8 @@
 python3 main_linear.py \
     --dataset imagenet100 \
     --backbone resnet18 \
-    --data_dir /datasets \
-    --train_dir imagenet-100/train \
-    --val_dir imagenet-100/val \
+    --train_data_path /datasets/imagenet-100/train \
+    --val_data_path /datasets/imagenet-100/val \
     --max_epochs 100 \
     --devices 0 \
     --accelerator gpu \
@@ -15,7 +14,7 @@ python3 main_linear.py \
     --weight_decay 0 \
     --batch_size 256 \
     --num_workers 4 \
-    --dali \
+    --data_format dali \
     --name byol-imagenet100-linear-eval \
     --pretrained_feature_extractor PATH \
     --project solo-learn \

bash_files/linear/imagenet-100/deepclusterv2_linear.sh → scripts/linear/imagenet-100/deepclusterv2_linear.sh

@@ -1,9 +1,8 @@
 python3 main_linear.py \
     --dataset imagenet100 \
     --backbone resnet18 \
-    --data_dir /data/datasets \
-    --train_dir imagenet-100/train \
-    --val_dir imagenet-100/val \
+    --train_data_path /datasets/imagenet-100/train \
+    --val_data_path /datasets/imagenet-100/val \
     --max_epochs 100 \
     --devices 0 \
     --accelerator gpu \
@@ -15,7 +14,7 @@ python3 main_linear.py \
     --weight_decay 0 \
     --batch_size 256 \
     --num_workers 5 \
-    --dali \
+    --data_format dali \
     --name deepclusterv2-imagenet100-linear-eval \
     --pretrained_feature_extractor PATH --project solo-learn \
     --entity unitn-mhug \

bash_files/linear/imagenet-100/dino_linear.sh → scripts/linear/imagenet-100/dino_linear.sh

@@ -1,9 +1,8 @@
 python3 main_linear.py \
     --dataset imagenet100 \
     --backbone resnet18 \
-    --data_dir /datasets \
-    --train_dir imagenet-100/train \
-    --val_dir imagenet-100/val \
+    --train_data_path /datasets/imagenet-100/train \
+    --val_data_path /datasets/imagenet-100/val \
     --max_epochs 100 \
     --devices 0 \
     --accelerator gpu \
@@ -15,7 +14,7 @@ python3 main_linear.py \
     --weight_decay 0 \
     --batch_size 256 \
     --num_workers 4 \
-    --dali \
+    --data_format dali \
     --name dino-imagenet100-linear-eval \
     --pretrained_feature_extractor PATH \
     --project solo-learn \

bash_files/linear/imagenet-100/general_linear.sh → scripts/linear/imagenet-100/general_linear.sh

@@ -1,9 +1,8 @@
 python3 main_linear.py \
     --dataset imagenet100 \
     --backbone resnet18 \
-    --data_dir /datasets \
-    --train_dir imagenet-100/train \
-    --val_dir imagenet-100/val \
+    --train_data_path /datasets/imagenet-100/train \
+    --val_data_path /datasets/imagenet-100/val \
     --max_epochs 100 \
     --devices 0,1 \
     --accelerator gpu \
@@ -17,7 +16,7 @@ python3 main_linear.py \
     --weight_decay 0 \
     --batch_size 128 \
     --num_workers 10 \
-    --dali \
+    --data_format dali \
     --name method-linear-eval \
     --pretrained_feature_extractor PATH \
     --project solo-learn \

bash_files/linear/imagenet-100/mocov2plus_linear.sh → scripts/linear/imagenet-100/mocov2plus_linear.sh

@@ -1,9 +1,8 @@
 python3 main_linear.py \
     --dataset imagenet100 \
     --backbone resnet18 \
-    --data_dir /datasets \
-    --train_dir imagenet-100/train \
-    --val_dir imagenet-100/val \
+    --train_data_path /datasets/imagenet-100/train \
+    --val_data_path /datasets/imagenet-100/val \
     --max_epochs 100 \
     --devices 0 \
     --accelerator gpu \
@@ -15,7 +14,7 @@ python3 main_linear.py \
     --weight_decay 0 \
     --batch_size 256 \
     --num_workers 10 \
-    --dali \
+    --data_format dali \
     --name mocov2plus-imagenet100-linear-eval \
     --pretrained_feature_extractor PATH \
     --project solo-learn \