Conditioning strip simsiam simclr

README_conditioning.md

@@ -0,0 +1,73 @@
+# Improving Transferability of Representations via Augmentation-Aware Self-Supervision
+
+Accepted to NeurIPS 2021
+
+<p align="center">
+<img width="762" alt="thumbnail" src="https://user-images.githubusercontent.com/4075389/138967888-29208bbe-d9e7-4bc7-b0b6-15ecbd5d277c.png">
+</p>
+
+**TL;DR:** Learning augmentation-aware information by conditioning on the encodings of two augmentations improves the transferability of representations. This is extention of AugSelf repo!
+
+## Dependencies
+
+```bash
+conda create -n AugSelf python=3.8 pytorch=1.7.1 torchvision=0.8.2 cudatoolkit=10.1 ignite -c pytorch
+conda activate AugSelf
+pip install scipy tensorboard kornia==0.4.1 sklearn
+
+conda create -n AugSelfConidtioning python=3.8
+conda activate AugSelfConditioning
+conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=10.1 -c pytorch
+pip install scipy tensorboard kornia==0.4.1 sklearn
+conda install -c conda-forge packaging
+conda install -c conda-forge wandb
+conda install ignite -c pytorch
+```
+
+## Checkpoints
+
+We provide ImageNet100-pretrained models in [this Dropbox link](https://www.dropbox.com/sh/0hjts19ysxebmaa/AABB6bF3QQWdIOCh9vocwTGGa?dl=0).
+
+## Pretraining
+
+We here provide SimSiam+ConditioningMLP pretraining scripts. For training the baseline (i.e., no MLP component), remove `--ss-crop` and `--ss-color` options. For using other frameworks like SimCLR, use the `--framework` option.
+
+### STL-10
+```
+script ovh_train_[***]_stl.sh
+```
+
+### ImageNet100
+
+```
+script ovh_train_[***]_imnet.sh
+```
+
+## Evaluation
+
+```
+script ovh_eval.sh
+```
+
+### linear evaluation
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python transfer_linear_eval.py \
+    --pretrain-data imagenet100 \
+    --ckpt CKPT \
+    --model resnet50 \
+    --dataset cifar10 \
+    --datadir DATADIR \
+    --metric top1
+```
+
+### few-shot
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python transfer_few_shot.py \
+    --pretrain-data imagenet100 \
+    --ckpt CKPT \
+    --model resnet50 \
+    --dataset cub200 \
+    --datadir DATADIR
+```

cond_utils.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import nn as nn
 
@@ -10,7 +12,8 @@
     "blur": 1,
     # "rot": 4,
     # "sol": 1,
-    "grayscale": 1
+    "grayscale": 1,
+    "color_diff": 3
 }
 
 
@@ -43,7 +46,7 @@ class AUG_DESC_TYPES:
 
 class AugProjector(nn.Module):
     def __init__(
-            self, args, proj_out_dim: int, proj_depth: int = 2
+            self, args, proj_out_dim: int, proj_depth: int = 2, proj_hidden_dim: Optional[int] = None, projector_last_bn: bool = False, projector_last_bn_affine: bool = True,
     ):
         super().__init__()
         self.num_backbone_features = args.num_backbone_features
@@ -54,6 +57,7 @@ def __init__(
         self.aug_cond = args.aug_cond or []
         self.aug_subset_sizes = {k: v for (k, v) in AUG_DESC_SIZE_CONFIG.items() if k in self.aug_cond}
         self.aug_inj_type = args.aug_inj_type
+        self.projector_last_bn = projector_last_bn
 
         print("Projector aug strategy:", self.aug_treatment)
         print("Conditioning projector on augmentations:", self.aug_subset_sizes)
@@ -147,11 +151,13 @@ def __init__(
             )
             self.projector = load_mlp(
                 projector_in,
-                args.num_backbone_features,
+                proj_hidden_dim or args.num_backbone_features,
                 proj_out_dim,
                 num_layers=proj_depth,
-                last_bn=False
+                last_bn=projector_last_bn,
+                last_bn_affine=projector_last_bn_affine,
             )
+            print(self.projector)
 
     def forward(self, x: torch.Tensor, aug_desc: torch.Tensor):
 
@@ -160,7 +166,7 @@ def forward(self, x: torch.Tensor, aug_desc: torch.Tensor):
 
             # print(f"pre {x.shape=}, {aug_desc.shape=}")
             if self.aug_inj_type == AUG_INJECTION_TYPES.proj_cat:
-                x = torch.concat([x, aug_desc], dim=1)
+                x = torch.cat([x, aug_desc], dim=1)
 
             elif self.aug_inj_type == AUG_INJECTION_TYPES.proj_add:
                 assert aug_desc.shape == x.shape, (x.shape, aug_desc.shape)