Upload 98 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+figs/in-loss-reg.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Wele Gedara Chaminda Bandara
+
+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 the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MODEL_ZOO.md

@@ -0,0 +1,20 @@
+The following links provide pre-trained models:
+# ResNet-18 Pre-trained Models
+| Dataset        |  d   | Lambda_BT | Lambda_Reg | Path to Pretrained Model | KNN Acc. | Linear Acc. |
+| ----------     | ---  | ---------- | ---------- | ------------------------ | -------- | ----------- |
+| CIFAR-10       | 1024 | 0.0078125  | 4.0        | 4wdhbpcf_0.0078125_1024_256_cifar10_model.pth     | 90.52    | 92.58        |
+| CIFAR-100      | 1024 | 0.0078125  | 4.0        | 76kk7scz_0.0078125_1024_256_cifar100_model.pth     | 61.25     | 69.31        |
+| TinyImageNet   | 1024 | 0.0009765  | 4.0        | 02azq6fs_0.0009765_1024_256_tiny_imagenet_model.pth     | 38.11    | 51.67        |
+| STL-10         | 1024 | 0.0078125  | 2.0        | i7det4xq_0.0078125_1024_256_stl10_model.pth     | 88.94     | 91.02        |
+
+
+
+# ResNet-50 Pre-trained Models
+| Dataset        |  d   | Lambda_BT | Lambda_Reg | Path to Pretrained Model | KNN Acc. | Linear Acc. |
+| ----------     | ---  | ---------- | ---------- | ------------------------ | -------- | ----------- |
+| CIFAR-10       | 1024 | 0.0078125  | 4.0        | v3gwgusq_0.0078125_1024_256_cifar10_model.pth     | 91.39     | 93.89        |
+| CIFAR-100      | 1024 | 0.0078125  | 4.0        | z6ngefw7_0.0078125_1024_256_cifar100_model_2000.pth     | 64.32     | 72.51        |
+| TinyImageNet   | 1024 | 0.0009765  | 4.0        | kxlkigsv_0.0009765_1024_256_tiny_imagenet_model_2000.pth     | 42.21     | 51.84        |
+| STL-10         | 1024 | 0.0078125  | 2.0        | pbknx38b_0.0078125_1024_256_stl10_model.pth     | 87.79     | 91.70        |
+| ImageNet       | 1024 | 0.0051  | 0.1        | 13awtq23_0.0051_8192_1024_imagenet_0.1_resnet50.pth     | -     | 72.1        |
+

README.md

@@ -1,6 +1,3 @@
----
-license: mit
----
 # Mixed Barlow Twins
 [**Guarding Barlow Twins Against Overfitting with Mixed Samples**](https://arxiv.org/abs/2312.02151)<br>
 

augmentations/augmentations_cifar.py

@@ -0,0 +1,190 @@
+# Copyright 2019 Google LLC
+#
+# 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
+#
+#     https://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.
+# ==============================================================================
+"""Base augmentations operators."""
+
+import numpy as np
+from PIL import Image, ImageOps, ImageEnhance
+
+# ImageNet code should change this value
+IMAGE_SIZE = 32
+import torch
+from torchvision import transforms
+
+
+def int_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval .
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled to
+      level/PARAMETER_MAX.
+
+  Returns:
+    An int that results from scaling `maxval` according to `level`.
+  """
+  return int(level * maxval / 10)
+
+
+def float_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval.
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled to
+      level/PARAMETER_MAX.
+
+  Returns:
+    A float that results from scaling `maxval` according to `level`.
+  """
+  return float(level) * maxval / 10.
+
+
+def sample_level(n):
+  return np.random.uniform(low=0.1, high=n)
+
+
+def autocontrast(pil_img, _):
+  return ImageOps.autocontrast(pil_img)
+
+
+def equalize(pil_img, _):
+  return ImageOps.equalize(pil_img)
+
+
+def posterize(pil_img, level):
+  level = int_parameter(sample_level(level), 4)
+  return ImageOps.posterize(pil_img, 4 - level)
+
+
+def rotate(pil_img, level):
+  degrees = int_parameter(sample_level(level), 30)
+  if np.random.uniform() > 0.5:
+    degrees = -degrees
+  return pil_img.rotate(degrees, resample=Image.BILINEAR)
+
+
+def solarize(pil_img, level):
+  level = int_parameter(sample_level(level), 256)
+  return ImageOps.solarize(pil_img, 256 - level)
+
+
+def shear_x(pil_img, level):
+  level = float_parameter(sample_level(level), 0.3)
+  if np.random.uniform() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, level, 0, 0, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def shear_y(pil_img, level):
+  level = float_parameter(sample_level(level), 0.3)
+  if np.random.uniform() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, 0, level, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def translate_x(pil_img, level):
+  level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
+  if np.random.random() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, level, 0, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def translate_y(pil_img, level):
+  level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
+  if np.random.random() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, 0, 0, 1, level),
+                           resample=Image.BILINEAR)
+
+
+# operation that overlaps with ImageNet-C's test set
+def color(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Color(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def contrast(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Contrast(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def brightness(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Brightness(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def sharpness(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Sharpness(pil_img).enhance(level)
+
+def random_resized_crop(pil_img, level):
+  return transforms.RandomResizedCrop(32)(pil_img)
+
+def random_flip(pil_img, level):
+  return transforms.RandomHorizontalFlip(p=0.5)(pil_img)
+
+def grayscale(pil_img, level):
+  return transforms.Grayscale(num_output_channels=3)(pil_img)
+
+augmentations = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y, grayscale #random_resized_crop, random_flip
+]
+
+augmentations_all = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y, color, contrast, brightness, sharpness, grayscale #, random_resized_crop, random_flip
+]
+
+def aug_cifar(image, preprocess, mixture_width=3, mixture_depth=-1, aug_severity=3):
+  """Perform AugMix augmentations and compute mixture.
+
+  Args:
+    image: PIL.Image input image
+    preprocess: Preprocessing function which should return a torch tensor.
+
+  Returns:
+    mixed: Augmented and mixed image.
+  """
+  aug_list = augmentations_all
+  # if args.all_ops:
+  #   aug_list = augmentations.augmentations_all
+
+  ws = np.float32(np.random.dirichlet([1] * mixture_width))
+  m = np.float32(np.random.beta(1, 1))
+
+  mix = torch.zeros_like(preprocess(image))
+  for i in range(mixture_width):
+    image_aug = image.copy()
+    depth = mixture_depth if mixture_depth > 0 else np.random.randint(
+        1, 4)
+    for _ in range(depth):
+      op = np.random.choice(aug_list)
+      image_aug = op(image_aug, aug_severity)
+    # Preprocessing commutes since all coefficients are convex
+    mix += ws[i] * preprocess(image_aug)
+
+  # mixed = (1 - m) * preprocess(image) + m * mix
+  return mix

augmentations/augmentations_stl.py

@@ -0,0 +1,190 @@
+# Copyright 2019 Google LLC
+#
+# 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
+#
+#     https://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.
+# ==============================================================================
+"""Base augmentations operators."""
+
+import numpy as np
+from PIL import Image, ImageOps, ImageEnhance
+
+# ImageNet code should change this value
+IMAGE_SIZE = 64
+import torch
+from torchvision import transforms
+
+
+def int_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval .
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled to
+      level/PARAMETER_MAX.
+
+  Returns:
+    An int that results from scaling `maxval` according to `level`.
+  """
+  return int(level * maxval / 10)
+
+
+def float_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval.
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled to
+      level/PARAMETER_MAX.
+
+  Returns:
+    A float that results from scaling `maxval` according to `level`.
+  """
+  return float(level) * maxval / 10.
+
+
+def sample_level(n):
+  return np.random.uniform(low=0.1, high=n)
+
+
+def autocontrast(pil_img, _):
+  return ImageOps.autocontrast(pil_img)
+
+
+def equalize(pil_img, _):
+  return ImageOps.equalize(pil_img)
+
+
+def posterize(pil_img, level):
+  level = int_parameter(sample_level(level), 4)
+  return ImageOps.posterize(pil_img, 4 - level)
+
+
+def rotate(pil_img, level):
+  degrees = int_parameter(sample_level(level), 30)
+  if np.random.uniform() > 0.5:
+    degrees = -degrees
+  return pil_img.rotate(degrees, resample=Image.BILINEAR)
+
+
+def solarize(pil_img, level):
+  level = int_parameter(sample_level(level), 256)
+  return ImageOps.solarize(pil_img, 256 - level)
+
+
+def shear_x(pil_img, level):
+  level = float_parameter(sample_level(level), 0.3)
+  if np.random.uniform() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, level, 0, 0, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def shear_y(pil_img, level):
+  level = float_parameter(sample_level(level), 0.3)
+  if np.random.uniform() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, 0, level, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def translate_x(pil_img, level):
+  level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
+  if np.random.random() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, level, 0, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def translate_y(pil_img, level):
+  level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
+  if np.random.random() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, 0, 0, 1, level),
+                           resample=Image.BILINEAR)
+
+
+# operation that overlaps with ImageNet-C's test set
+def color(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Color(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def contrast(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Contrast(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def brightness(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Brightness(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def sharpness(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Sharpness(pil_img).enhance(level)
+
+def random_resized_crop(pil_img, level):
+  return transforms.RandomResizedCrop(32)(pil_img)
+
+def random_flip(pil_img, level):
+  return transforms.RandomHorizontalFlip(p=0.5)(pil_img)
+
+def grayscale(pil_img, level):
+  return transforms.Grayscale(num_output_channels=3)(pil_img)
+
+augmentations = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y, grayscale #random_resized_crop, random_flip
+]
+
+augmentations_all = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y, color, contrast, brightness, sharpness, grayscale #, random_resized_crop, random_flip
+]
+
+def aug_stl(image, preprocess, mixture_width=3, mixture_depth=-1, aug_severity=3):
+  """Perform AugMix augmentations and compute mixture.
+
+  Args:
+    image: PIL.Image input image
+    preprocess: Preprocessing function which should return a torch tensor.
+
+  Returns:
+    mixed: Augmented and mixed image.
+  """
+  aug_list = augmentations
+  # if args.all_ops:
+  #   aug_list = augmentations.augmentations_all
+
+  ws = np.float32(np.random.dirichlet([1] * mixture_width))
+  m = np.float32(np.random.beta(1, 1))
+
+  mix = torch.zeros_like(preprocess(image))
+  for i in range(mixture_width):
+    image_aug = image.copy()
+    depth = mixture_depth if mixture_depth > 0 else np.random.randint(
+        1, 4)
+    for _ in range(depth):
+      op = np.random.choice(aug_list)
+      image_aug = op(image_aug, aug_severity)
+    # Preprocessing commutes since all coefficients are convex
+    mix += ws[i] * preprocess(image_aug)
+
+  mixed = (1 - m) * preprocess(image) + m * mix
+  return mixed

augmentations/augmentations_tiny.py

@@ -0,0 +1,190 @@
+# Copyright 2019 Google LLC
+#
+# 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
+#
+#     https://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.
+# ==============================================================================
+"""Base augmentations operators."""
+
+import numpy as np
+from PIL import Image, ImageOps, ImageEnhance
+
+# ImageNet code should change this value
+IMAGE_SIZE = 64
+import torch
+from torchvision import transforms
+
+
+def int_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval .
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled to
+      level/PARAMETER_MAX.
+
+  Returns:
+    An int that results from scaling `maxval` according to `level`.
+  """
+  return int(level * maxval / 10)
+
+
+def float_parameter(level, maxval):
+  """Helper function to scale `val` between 0 and maxval.
+
+  Args:
+    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
+    maxval: Maximum value that the operation can have. This will be scaled to
+      level/PARAMETER_MAX.
+
+  Returns:
+    A float that results from scaling `maxval` according to `level`.
+  """
+  return float(level) * maxval / 10.
+
+
+def sample_level(n):
+  return np.random.uniform(low=0.1, high=n)
+
+
+def autocontrast(pil_img, _):
+  return ImageOps.autocontrast(pil_img)
+
+
+def equalize(pil_img, _):
+  return ImageOps.equalize(pil_img)
+
+
+def posterize(pil_img, level):
+  level = int_parameter(sample_level(level), 4)
+  return ImageOps.posterize(pil_img, 4 - level)
+
+
+def rotate(pil_img, level):
+  degrees = int_parameter(sample_level(level), 30)
+  if np.random.uniform() > 0.5:
+    degrees = -degrees
+  return pil_img.rotate(degrees, resample=Image.BILINEAR)
+
+
+def solarize(pil_img, level):
+  level = int_parameter(sample_level(level), 256)
+  return ImageOps.solarize(pil_img, 256 - level)
+
+
+def shear_x(pil_img, level):
+  level = float_parameter(sample_level(level), 0.3)
+  if np.random.uniform() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, level, 0, 0, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def shear_y(pil_img, level):
+  level = float_parameter(sample_level(level), 0.3)
+  if np.random.uniform() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, 0, level, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def translate_x(pil_img, level):
+  level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
+  if np.random.random() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, level, 0, 1, 0),
+                           resample=Image.BILINEAR)
+
+
+def translate_y(pil_img, level):
+  level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
+  if np.random.random() > 0.5:
+    level = -level
+  return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
+                           Image.AFFINE, (1, 0, 0, 0, 1, level),
+                           resample=Image.BILINEAR)
+
+
+# operation that overlaps with ImageNet-C's test set
+def color(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Color(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def contrast(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Contrast(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def brightness(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Brightness(pil_img).enhance(level)
+
+
+# operation that overlaps with ImageNet-C's test set
+def sharpness(pil_img, level):
+    level = float_parameter(sample_level(level), 1.8) + 0.1
+    return ImageEnhance.Sharpness(pil_img).enhance(level)
+
+def random_resized_crop(pil_img, level):
+  return transforms.RandomResizedCrop(32)(pil_img)
+
+def random_flip(pil_img, level):
+  return transforms.RandomHorizontalFlip(p=0.5)(pil_img)
+
+def grayscale(pil_img, level):
+  return transforms.Grayscale(num_output_channels=3)(pil_img)
+
+augmentations = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y, grayscale #random_resized_crop, random_flip
+]
+
+augmentations_all = [
+    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
+    translate_x, translate_y, color, contrast, brightness, sharpness, grayscale #, random_resized_crop, random_flip
+]
+
+def aug_tiny(image, preprocess, mixture_width=3, mixture_depth=-1, aug_severity=3):
+  """Perform AugMix augmentations and compute mixture.
+
+  Args:
+    image: PIL.Image input image
+    preprocess: Preprocessing function which should return a torch tensor.
+
+  Returns:
+    mixed: Augmented and mixed image.
+  """
+  aug_list = augmentations
+  # if args.all_ops:
+  #   aug_list = augmentations.augmentations_all
+
+  ws = np.float32(np.random.dirichlet([1] * mixture_width))
+  m = np.float32(np.random.beta(1, 1))
+
+  mix = torch.zeros_like(preprocess(image))
+  for i in range(mixture_width):
+    image_aug = image.copy()
+    depth = mixture_depth if mixture_depth > 0 else np.random.randint(
+        1, 4)
+    for _ in range(depth):
+      op = np.random.choice(aug_list)
+      image_aug = op(image_aug, aug_severity)
+    # Preprocessing commutes since all coefficients are convex
+    mix += ws[i] * preprocess(image_aug)
+
+  mixed = (1 - m) * preprocess(image) + m * mix
+  return mixed

data_statistics.py

@@ -0,0 +1,61 @@
+def get_data_mean_and_stdev(dataset):
+    if dataset == 'CIFAR10' or dataset == 'CIFAR100':
+        mean = [0.5, 0.5, 0.5]
+        std  = [0.5, 0.5, 0.5]
+    elif dataset == 'STL-10':
+        mean = [0.491, 0.482, 0.447]
+        std  = [0.247, 0.244, 0.262]
+    elif dataset == 'ImageNet':
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+    elif dataset == 'aircraft':
+        mean = [0.486, 0.507, 0.525]
+        std  = [0.266, 0.260, 0.276]
+    elif dataset == 'cu_birds':
+        mean = [0.483, 0.491, 0.424] 
+        std  = [0.228, 0.224, 0.259]
+    elif dataset == 'dtd':
+        mean = [0.533, 0.474, 0.426]
+        std  = [0.261, 0.250, 0.259]
+    elif dataset == 'fashionmnist':
+        mean = [0.348, 0.348, 0.348] 
+        std  = [0.347, 0.347, 0.347]
+    elif dataset == 'mnist':
+        mean = [0.170, 0.170, 0.170]
+        std  = [0.320, 0.320, 0.320]
+    elif dataset == 'traffic_sign':
+        mean = [0.335, 0.291, 0.295]
+        std  = [0.267, 0.249, 0.251]
+    elif dataset == 'vgg_flower':
+        mean = [0.518, 0.410, 0.329]
+        std  = [0.296, 0.249, 0.285]
+    else:
+        raise Exception('Dataset %s not supported.'%dataset)
+    return mean, std
+
+def get_data_nclass(dataset):
+    if dataset == 'cifar10':
+        nclass = 10
+    elif dataset == 'cifar100cifar10':
+        nclass = 100
+    elif dataset == 'stl-10':
+        nclass = 10
+    elif dataset == 'ImageNet':
+        nclass = 1000
+    elif dataset == 'aircraft':
+        nclass = 102
+    elif dataset == 'cu_birds':
+        nclass = 200
+    elif dataset == 'dtd':
+        nclass = 47
+    elif dataset == 'fashionmnist':
+        nclass = 10
+    elif dataset == 'mnist':
+        nclass = 10
+    elif dataset == 'traffic_sign':
+        nclass = 43
+    elif dataset == 'vgg_flower':
+        nclass = 102
+    else:
+        raise Exception('Dataset %s not supported.'%dataset)
+    return nclass

download_imagenet.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+# https://gist.github.com/BIGBALLON/8a71d225eff18d88e469e6ea9b39cef4
+cd /mnt/store/wbandar1/datasets
+wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar --no-check-certificate
+wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar --no-check-certificate
+
+#
+# script to extract ImageNet dataset
+# ILSVRC2012_img_train.tar (about 138 GB)
+# ILSVRC2012_img_val.tar (about 6.3 GB)
+# make sure ILSVRC2012_img_train.tar & ILSVRC2012_img_val.tar in your current directory
+#
+#  https://github.com/facebook/fb.resnet.torch/blob/master/INSTALL.md
+# 
+#  train/
+#  ├── n01440764
+#  │   ├── n01440764_10026.JPEG
+#  │   ├── n01440764_10027.JPEG
+#  │   ├── ......
+#  ├── ......
+#  val/
+#  ├── n01440764
+#  │   ├── ILSVRC2012_val_00000293.JPEG
+#  │   ├── ILSVRC2012_val_00002138.JPEG
+#  │   ├── ......
+#  ├── ......
+#
+#
+# Extract the training data:
+#
+mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
+tar -xvf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
+find . -name "*.tar" | while read NAME ; do mkdir -p "${NAME%.tar}"; tar -xvf "${NAME}" -C "${NAME%.tar}"; rm -f "${NAME}"; done
+cd ..
+#
+# Extract the validation data and move images to subfolders:
+#
+mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xvf ILSVRC2012_img_val.tar
+wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash
+#
+# Check total files after extract
+#
+#  $ find train/ -name "*.JPEG" | wc -l
+#  1281167
+#  $ find val/ -name "*.JPEG" | wc -l
+#  50000
+#

environment.yml

@@ -0,0 +1,188 @@
+name: ssl-aug
+channels:
+  - pytorch
+  - anaconda
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - blas=1.0=mkl
+  - bottleneck=1.3.4=py38hce1f21e_0
+  - brotlipy=0.7.0=py38h27cfd23_1003
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2022.6.15=ha878542_0
+  - cairo=1.16.0=hcf35c78_1003
+  - certifi=2022.6.15=py38h578d9bd_0
+  - cffi=1.15.0=py38h7f8727e_0
+  - charset-normalizer=2.0.4=pyhd3eb1b0_0
+  - cryptography=37.0.1=py38h9ce1e76_0
+  - cudatoolkit=11.3.1=h2bc3f7f_2
+  - dataclasses=0.8=pyh6d0b6a4_7
+  - dbus=1.13.18=hb2f20db_0
+  - expat=2.4.8=h27087fc_0
+  - ffmpeg=4.3.2=hca11adc_0
+  - fontconfig=2.14.0=h8e229c2_0
+  - freetype=2.11.0=h70c0345_0
+  - fvcore=0.1.5.post20220512=pyhd8ed1ab_0
+  - gettext=0.19.8.1=hd7bead4_3
+  - gh=2.12.1=ha8f183a_0
+  - giflib=5.2.1=h7b6447c_0
+  - glib=2.66.3=h58526e2_0
+  - gmp=6.2.1=h295c915_3
+  - gnutls=3.6.15=he1e5248_0
+  - graphite2=1.3.14=h295c915_1
+  - gst-plugins-base=1.14.5=h0935bb2_2
+  - gstreamer=1.14.5=h36ae1b5_2
+  - harfbuzz=2.4.0=h9f30f68_3
+  - hdf5=1.10.6=hb1b8bf9_0
+  - icu=64.2=he1b5a44_1
+  - idna=3.3=pyhd3eb1b0_0
+  - intel-openmp=2021.4.0=h06a4308_3561
+  - iopath=0.1.9=pyhd8ed1ab_0
+  - jasper=1.900.1=hd497a04_4
+  - jpeg=9e=h7f8727e_0
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - libblas=3.9.0=12_linux64_mkl
+  - libcblas=3.9.0=12_linux64_mkl
+  - libclang=9.0.1=default_hb4e5071_5
+  - libedit=3.1.20210910=h7f8727e_0
+  - libffi=3.2.1=hf484d3e_1007
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgfortran-ng=7.5.0=ha8ba4b0_17
+  - libgfortran4=7.5.0=ha8ba4b0_17
+  - libglib=2.66.3=hbe7bbb4_0
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.16=h7f8727e_2
+  - libidn2=2.3.2=h7f8727e_0
+  - liblapack=3.9.0=12_linux64_mkl
+  - liblapacke=3.9.0=12_linux64_mkl
+  - libllvm9=9.0.1=h4a3c616_1
+  - libopencv=4.5.1=py38h703c3c0_0
+  - libpng=1.6.37=hbc83047_0
+  - libprotobuf=3.15.8=h780b84a_1
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtasn1=4.16.0=h27cfd23_0
+  - libtiff=4.2.0=h2818925_1
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=2.32.1=h7f98852_1000
+  - libuv=1.40.0=h7b6447c_0
+  - libwebp=1.2.2=h55f646e_0
+  - libwebp-base=1.2.2=h7f8727e_0
+  - libxcb=1.15=h7f8727e_0
+  - libxkbcommon=0.10.0=he1b5a44_0
+  - libxml2=2.9.9=hea5a465_1
+  - lz4-c=1.9.3=h295c915_1
+  - mkl=2021.4.0=h06a4308_640
+  - mkl-service=2.4.0=py38h7f8727e_0
+  - mkl_fft=1.3.1=py38hd3c417c_0
+  - mkl_random=1.2.2=py38h51133e4_0
+  - ncurses=6.3=h7f8727e_2
+  - nettle=3.7.3=hbbd107a_1
+  - nspr=4.33=h295c915_0
+  - nss=3.46.1=hab99668_0
+  - numexpr=2.8.1=py38h6abb31d_0
+  - numpy=1.22.3=py38he7a7128_0
+  - numpy-base=1.22.3=py38hf524024_0
+  - opencv=4.5.1=py38h578d9bd_0
+  - openh264=2.1.1=h4ff587b_0
+  - openssl=1.1.1o=h166bdaf_0
+  - packaging=21.3=pyhd3eb1b0_0
+  - pandas=1.4.2=py38h295c915_0
+  - pcre=8.45=h295c915_0
+  - pillow=9.0.1=py38h22f2fdc_0
+  - pip=21.2.4=py38h06a4308_0
+  - pixman=0.38.0=h7b6447c_0
+  - portalocker=2.3.0=py38h06a4308_0
+  - protobuf=3.15.8=py38h709712a_0
+  - py-opencv=4.5.1=py38h81c977d_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyopenssl=22.0.0=pyhd3eb1b0_0
+  - pyparsing=3.0.9=pyhd8ed1ab_0
+  - pysocks=1.7.1=py38h06a4308_0
+  - python=3.8.0=h0371630_2
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - python_abi=3.8=2_cp38
+  - pytorch=1.11.0=py3.8_cuda11.3_cudnn8.2.0_0
+  - pytorch-mutex=1.0=cuda
+  - pytz=2021.3=pyhd3eb1b0_0
+  - pyyaml=6.0=py38h7f8727e_1
+  - qt=5.12.5=hd8c4c69_1
+  - readline=7.0=h7b6447c_5
+  - requests=2.27.1=pyhd3eb1b0_0
+  - setuptools=61.2.0=py38h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - sqlite=3.33.0=h62c20be_0
+  - tabulate=0.8.9=py38h06a4308_0
+  - tensorboardx=2.5.1=pyhd8ed1ab_0
+  - termcolor=1.1.0=py38h06a4308_1
+  - tk=8.6.12=h1ccaba5_0
+  - torchvision=0.12.0=py38_cu113
+  - tqdm=4.64.0=py38h06a4308_0
+  - typing_extensions=4.1.1=pyh06a4308_0
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - x264=1!161.3030=h7f98852_1
+  - xorg-kbproto=1.0.7=h7f98852_1002
+  - xorg-libice=1.0.10=h7f98852_0
+  - xorg-libsm=1.2.3=hd9c2040_1000
+  - xorg-libx11=1.7.2=h7f98852_0
+  - xorg-libxext=1.3.4=h7f98852_1
+  - xorg-libxrender=0.9.10=h7f98852_1003
+  - xorg-renderproto=0.11.1=h7f98852_1002
+  - xorg-xextproto=7.3.0=h7f98852_1002
+  - xorg-xproto=7.0.31=h27cfd23_1007
+  - xz=5.2.5=h7f8727e_1
+  - yacs=0.1.6=pyhd3eb1b0_1
+  - yaml=0.2.5=h7b6447c_0
+  - zip=3.0=h7f98852_1
+  - zlib=1.2.12=h7f8727e_2
+  - zstd=1.5.2=ha4553b6_0
+  - pip:
+    - absl-py==1.1.0
+    - appdirs==1.4.4
+    - cachetools==5.2.0
+    - click==8.1.7
+    - contourpy==1.0.6
+    - cycler==0.11.0
+    - decord==0.6.0
+    - deepspeed==0.5.8
+    - docker-pycreds==0.4.0
+    - einops==0.4.1
+    - filelock==3.7.1
+    - fonttools==4.38.0
+    - future==0.18.2
+    - gitdb==4.0.10
+    - gitpython==3.1.33
+    - google-auth==2.7.0
+    - google-auth-oauthlib==0.4.6
+    - grpcio==1.46.3
+    - hjson==3.0.2
+    - imageio==2.22.2
+    - importlib-metadata==4.11.4
+    - kiwisolver==1.4.4
+    - markdown==3.3.7
+    - matplotlib==3.6.1
+    - ninja==1.10.2.3
+    - oauthlib==3.2.0
+    - pathtools==0.1.2
+    - psutil==5.9.1
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - requests-oauthlib==1.3.1
+    - rsa==4.8
+    - scipy==1.9.0
+    - sentry-sdk==1.30.0
+    - setproctitle==1.3.2
+    - smmap==5.0.0
+    - tensorboard==2.9.1
+    - tensorboard-data-server==0.6.1
+    - tensorboard-plugin-wit==1.8.1
+    - thop==0.1.1-2209072238
+    - timm==0.4.12
+    - triton==1.1.1
+    - urllib3==1.26.16
+    - wandb==0.15.9
+    - werkzeug==2.1.2
+    - zipp==3.8.0
+prefix: /home/wbandar1/anaconda3/envs/ssl-aug

evaluate_imagenet.py

@@ -0,0 +1,289 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from pathlib import Path
+import argparse
+import json
+import os
+import random
+import signal
+import sys
+import time
+import urllib
+
+from torch import nn, optim
+from torchvision import models, datasets, transforms
+import torch
+import torchvision
+import wandb
+
+parser = argparse.ArgumentParser(description='Evaluate resnet50 features on ImageNet')
+parser.add_argument('data', type=Path, metavar='DIR',
+                    help='path to dataset')
+parser.add_argument('pretrained', type=Path, metavar='FILE',
+                    help='path to pretrained model')
+parser.add_argument('--weights', default='freeze', type=str,
+                    choices=('finetune', 'freeze'),
+                    help='finetune or freeze resnet weights')
+parser.add_argument('--train-percent', default=100, type=int,
+                    choices=(100, 10, 1),
+                    help='size of traing set in percent')
+parser.add_argument('--workers', default=8, type=int, metavar='N',
+                    help='number of data loader workers')
+parser.add_argument('--epochs', default=100, type=int, metavar='N',
+                    help='number of total epochs to run')
+parser.add_argument('--batch-size', default=256, type=int, metavar='N',
+                    help='mini-batch size')
+parser.add_argument('--lr-backbone', default=0.0, type=float, metavar='LR',
+                    help='backbone base learning rate')
+parser.add_argument('--lr-classifier', default=0.3, type=float, metavar='LR',
+                    help='classifier base learning rate')
+parser.add_argument('--weight-decay', default=1e-6, type=float, metavar='W',
+                    help='weight decay')
+parser.add_argument('--print-freq', default=100, type=int, metavar='N',
+                    help='print frequency')
+parser.add_argument('--checkpoint-dir', default='/mnt/store/wbandar1/projects/ssl-aug-artifacts/', type=Path,
+                    metavar='DIR', help='path to checkpoint directory')
+
+
+def main():
+    args = parser.parse_args()
+    if args.train_percent in {1, 10}:
+        args.train_files = urllib.request.urlopen(f'https://raw.githubusercontent.com/google-research/simclr/master/imagenet_subsets/{args.train_percent}percent.txt').readlines()
+    args.ngpus_per_node = torch.cuda.device_count()
+    if 'SLURM_JOB_ID' in os.environ:
+        signal.signal(signal.SIGUSR1, handle_sigusr1)
+        signal.signal(signal.SIGTERM, handle_sigterm)
+    # single-node distributed training
+    args.rank = 0
+    args.dist_url = f'tcp://localhost:{random.randrange(49152, 65535)}'
+    args.world_size = args.ngpus_per_node
+    torch.multiprocessing.spawn(main_worker, (args,), args.ngpus_per_node)
+
+
+def main_worker(gpu, args):
+    args.rank += gpu
+    torch.distributed.init_process_group(
+        backend='nccl', init_method=args.dist_url,
+        world_size=args.world_size, rank=args.rank)
+
+    # initializing wandb
+    if args.rank == 0:
+        run = wandb.init(project="bt-in1k-eval", config=args, dir='/mnt/store/wbandar1/projects/ssl-aug-artifacts/wandb_logs/')
+        run_id = wandb.run.id
+        args.checkpoint_dir=Path(os.path.join(args.checkpoint_dir, run_id))
+
+    if args.rank == 0:
+        args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        stats_file = open(args.checkpoint_dir / 'stats.txt', 'a', buffering=1)
+        print(' '.join(sys.argv))
+        print(' '.join(sys.argv), file=stats_file)
+
+    torch.cuda.set_device(gpu)
+    torch.backends.cudnn.benchmark = True
+
+    model = models.resnet50().cuda(gpu)
+    state_dict = torch.load(args.pretrained, map_location='cpu')
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    assert missing_keys == ['fc.weight', 'fc.bias'] and unexpected_keys == []
+    model.fc.weight.data.normal_(mean=0.0, std=0.01)
+    model.fc.bias.data.zero_()
+    if args.weights == 'freeze':
+        model.requires_grad_(False)
+        model.fc.requires_grad_(True)
+    classifier_parameters, model_parameters = [], []
+    for name, param in model.named_parameters():
+        if name in {'fc.weight', 'fc.bias'}:
+            classifier_parameters.append(param)
+        else:
+            model_parameters.append(param)
+    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
+
+    criterion = nn.CrossEntropyLoss().cuda(gpu)
+
+    param_groups = [dict(params=classifier_parameters, lr=args.lr_classifier)]
+    if args.weights == 'finetune':
+        param_groups.append(dict(params=model_parameters, lr=args.lr_backbone))
+    optimizer = optim.SGD(param_groups, 0, momentum=0.9, weight_decay=args.weight_decay)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs)
+
+    # automatically resume from checkpoint if it exists
+    if (args.checkpoint_dir / 'checkpoint.pth').is_file():
+        ckpt = torch.load(args.checkpoint_dir / 'checkpoint.pth',
+                          map_location='cpu')
+        start_epoch = ckpt['epoch']
+        best_acc = ckpt['best_acc']
+        model.load_state_dict(ckpt['model'])
+        optimizer.load_state_dict(ckpt['optimizer'])
+        scheduler.load_state_dict(ckpt['scheduler'])
+    else:
+        start_epoch = 0
+        best_acc = argparse.Namespace(top1=0, top5=0)
+
+    # Data loading code
+    traindir = args.data / 'train'
+    valdir = args.data / 'val'
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+
+    train_dataset = datasets.ImageFolder(traindir, transforms.Compose([
+            transforms.RandomResizedCrop(224),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            normalize,
+        ]))
+    val_dataset = datasets.ImageFolder(valdir, transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            normalize,
+        ]))
+
+    if args.train_percent in {1, 10}:
+        train_dataset.samples = []
+        for fname in args.train_files:
+            fname = fname.decode().strip()
+            cls = fname.split('_')[0]
+            train_dataset.samples.append(
+                (traindir / cls / fname, train_dataset.class_to_idx[cls]))
+
+    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
+    kwargs = dict(batch_size=args.batch_size // args.world_size, num_workers=args.workers, pin_memory=True)
+    train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, **kwargs)
+    val_loader = torch.utils.data.DataLoader(val_dataset, **kwargs)
+
+    start_time = time.time()
+    for epoch in range(start_epoch, args.epochs):
+        # train
+        if args.weights == 'finetune':
+            model.train()
+        elif args.weights == 'freeze':
+            model.eval()
+        else:
+            assert False
+        train_sampler.set_epoch(epoch)
+        for step, (images, target) in enumerate(train_loader, start=epoch * len(train_loader)):
+            output = model(images.cuda(gpu, non_blocking=True))
+            loss = criterion(output, target.cuda(gpu, non_blocking=True))
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if step % args.print_freq == 0:
+                torch.distributed.reduce(loss.div_(args.world_size), 0)
+                if args.rank == 0:
+                    pg = optimizer.param_groups
+                    lr_classifier = pg[0]['lr']
+                    lr_backbone = pg[1]['lr'] if len(pg) == 2 else 0
+                    stats = dict(epoch=epoch, step=step, lr_backbone=lr_backbone,
+                                 lr_classifier=lr_classifier, loss=loss.item(),
+                                 time=int(time.time() - start_time))
+                    print(json.dumps(stats))
+                    print(json.dumps(stats), file=stats_file)
+                    run.log(
+                        {
+                            "epoch": epoch,
+                            "step": step,
+                            "lr_backbone": lr_backbone,
+                            "lr_classifier": lr_classifier,
+                            "loss": loss.item(),
+                            "time": int(time.time() - start_time),
+                            }
+                            )
+
+        # evaluate
+        model.eval()
+        if args.rank == 0:
+            top1 = AverageMeter('Acc@1')
+            top5 = AverageMeter('Acc@5')
+            with torch.no_grad():
+                for images, target in val_loader:
+                    output = model(images.cuda(gpu, non_blocking=True))
+                    acc1, acc5 = accuracy(output, target.cuda(gpu, non_blocking=True), topk=(1, 5))
+                    top1.update(acc1[0].item(), images.size(0))
+                    top5.update(acc5[0].item(), images.size(0))
+            best_acc.top1 = max(best_acc.top1, top1.avg)
+            best_acc.top5 = max(best_acc.top5, top5.avg)
+            stats = dict(epoch=epoch, acc1=top1.avg, acc5=top5.avg, best_acc1=best_acc.top1, best_acc5=best_acc.top5)
+            print(json.dumps(stats))
+            print(json.dumps(stats), file=stats_file)
+            run.log(
+                {
+                    "epoch": epoch,
+                    "eval_acc1": top1.avg,
+                    "eval_acc5": top5.avg,
+                    "eval_best_acc1": best_acc.top1,
+                    "eval_best_acc5": best_acc.top5,
+                    }
+                )
+
+        # sanity check
+        if args.weights == 'freeze':
+            reference_state_dict = torch.load(args.pretrained, map_location='cpu')
+            model_state_dict = model.module.state_dict()
+            for k in reference_state_dict:
+                assert torch.equal(model_state_dict[k].cpu(), reference_state_dict[k]), k
+
+        scheduler.step()
+        if args.rank == 0:
+            state = dict(
+                epoch=epoch + 1, best_acc=best_acc, model=model.state_dict(),
+                optimizer=optimizer.state_dict(), scheduler=scheduler.state_dict())
+            torch.save(state, args.checkpoint_dir / 'checkpoint.pth')
+    wandb.finish()
+
+
+def handle_sigusr1(signum, frame):
+    os.system(f'scontrol requeue {os.getenv("SLURM_JOB_ID")}')
+    exit()
+
+
+def handle_sigterm(signum, frame):
+    pass
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self, name, fmt=':f'):
+        self.name = name
+        self.fmt = fmt
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
+        return fmtstr.format(**self.__dict__)
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
+            res.append(correct_k.mul_(100.0 / batch_size))
+        return res
+
+
+if __name__ == '__main__':
+    main()

evaluate_transfer.py

@@ -0,0 +1,168 @@
+import argparse
+
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from thop import profile, clever_format
+from torch.utils.data import DataLoader
+from transfer_datasets import TRANSFER_DATASET
+import torchvision.transforms as transforms
+from data_statistics import get_data_mean_and_stdev, get_data_nclass
+from tqdm import tqdm
+
+import utils
+
+import wandb
+
+import torchvision
+
+def load_transform(dataset, size=32):
+    mean, std = get_data_mean_and_stdev(dataset)
+    transform = transforms.Compose([
+        transforms.Resize((size, size)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=mean, std=std)])
+    return transform
+
+class Net(nn.Module):
+    def __init__(self, num_class, pretrained_path, dataset, arch):
+        super(Net, self).__init__()
+
+        if arch=='resnet18':
+            embedding_size = 512
+        elif arch=='resnet50':
+            embedding_size = 2048
+        else:
+            raise NotImplementedError
+
+        # encoder
+        from model import Model
+        self.f = Model(dataset=dataset, arch=arch).f
+        # classifier
+        self.fc = nn.Linear(embedding_size, num_class, bias=True)
+        self.load_state_dict(torch.load(pretrained_path, map_location='cpu'), strict=False)
+
+    def forward(self, x):
+        x = self.f(x)
+        feature = torch.flatten(x, start_dim=1)
+        out = self.fc(feature)
+        return out
+
+# train or test for one epoch
+def train_val(net, data_loader, train_optimizer):
+    is_train = train_optimizer is not None
+    net.train() if is_train else net.eval()
+
+    total_loss, total_correct_1, total_correct_5, total_num, data_bar = 0.0, 0.0, 0.0, 0, tqdm(data_loader)
+    with (torch.enable_grad() if is_train else torch.no_grad()):
+        for data, target in data_bar:
+            data, target = data.cuda(non_blocking=True), target.cuda(non_blocking=True)
+            out = net(data)
+            loss = loss_criterion(out, target)
+
+            if is_train:
+                train_optimizer.zero_grad()
+                loss.backward()
+                train_optimizer.step()
+
+            total_num += data.size(0)
+            total_loss += loss.item() * data.size(0)
+            prediction = torch.argsort(out, dim=-1, descending=True)
+            total_correct_1 += torch.sum((prediction[:, 0:1] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
+            total_correct_5 += torch.sum((prediction[:, 0:5] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
+
+            data_bar.set_description('{} Epoch: [{}/{}] Loss: {:.4f} ACC@1: {:.2f}% ACC@5: {:.2f}% model: {}'
+                                     .format('Train' if is_train else 'Test', epoch, epochs, total_loss / total_num,
+                                             total_correct_1 / total_num * 100, total_correct_5 / total_num * 100,
+                                             model_path.split('/')[-1]))
+
+    return total_loss / total_num, total_correct_1 / total_num * 100, total_correct_5 / total_num * 100
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Linear Evaluation')
+    parser.add_argument('--dataset', default='cifar10', type=str, help='Pre-trained dataset.', choices=['cifar10', 'cifar100', 'stl10', 'tiny_imagenet'])
+    parser.add_argument('--transfer_dataset', default='cifar10', type=str, help='Transfer dataset (i.e., testing dataset)', choices=['cifar10', 'cifar100', 'stl-10', 'aircraft', 'cu_birds', 'dtd', 'fashionmnist', 'mnist', 'traffic_sign', 'vgg_flower'])
+    parser.add_argument('--arch', default='resnet50', type=str, help='Backbone architecture for experiments', choices=['resnet50', 'resnet18'])
+    parser.add_argument('--model_path', type=str, default='results/Barlow_Twins/0.005_64_128_model.pth',
+                        help='The base string of the pretrained model path')
+    parser.add_argument('--batch_size', type=int, default=128, help='Number of images in each mini-batch')
+    parser.add_argument('--epochs', type=int, default=100, help='Number of sweeps over the dataset to train')
+    parser.add_argument('--screen', type=str, help='screen session id')
+    # wandb related args
+    parser.add_argument('--wandb_group', type=str, help='group for wandb')
+
+    args = parser.parse_args()
+    
+    wandb.init(project=f"Barlow-Twins-MixUp-TransferLearn-[{args.dataset}-to-X]-{args.arch}", config=args, dir='/data/wbandar1/projects/ssl-aug-artifacts/wandb_logs/', group=args.wandb_group, name=f'{args.transfer_dataset}')
+    run_id = wandb.run.id
+
+    model_path, batch_size, epochs = args.model_path, args.batch_size, args.epochs
+    dataset = args.dataset
+    transfer_dataset = args.transfer_dataset
+    
+    if dataset in ['cifar10', 'cifar100']:
+        print("reshaping data into 32x32")
+        resize = 32
+    else:
+        print("reshaping data into 64x64")
+        resize = 64
+        
+    train_data = TRANSFER_DATASET[args.transfer_dataset](train=True, image_transforms=load_transform(args.transfer_dataset, resize))
+    test_data = TRANSFER_DATASET[args.transfer_dataset](train=False, image_transforms=load_transform(args.transfer_dataset, resize))
+    
+    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True)
+    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)
+    
+    model = Net(num_class=get_data_nclass(args.transfer_dataset), pretrained_path=model_path, dataset=dataset, arch=args.arch).cuda()
+    for param in model.f.parameters():
+        param.requires_grad = False
+
+    # optimizer with lr sheduler
+    # lr_start, lr_end = 1e-2, 1e-6
+    # gamma = (lr_end / lr_start) ** (1 / epochs)
+    # optimizer = optim.Adam(model.fc.parameters(), lr=lr_start, weight_decay=5e-6)
+    # scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)
+
+    # adpoted from 
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
+    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [60, 80], gamma=0.1)
+    
+    # optimizer with no sheuduler
+    # optimizer = optim.Adam(model.fc.parameters(), lr=1e-3, weight_decay=1e-6)
+
+    loss_criterion = nn.CrossEntropyLoss()
+    results = {'train_loss': [], 'train_acc@1': [], 'train_acc@5': [],
+               'test_loss': [], 'test_acc@1': [], 'test_acc@5': []}
+
+    save_name = model_path.split('.pth')[0] + '_linear.csv'
+
+    best_acc = 0.0
+    for epoch in range(1, epochs + 1):
+        train_loss, train_acc_1, train_acc_5 = train_val(model, train_loader, optimizer)
+        results['train_loss'].append(train_loss)
+        results['train_acc@1'].append(train_acc_1)
+        results['train_acc@5'].append(train_acc_5)
+        test_loss, test_acc_1, test_acc_5 = train_val(model, test_loader, None)
+        results['test_loss'].append(test_loss)
+        results['test_acc@1'].append(test_acc_1)
+        results['test_acc@5'].append(test_acc_5)
+        # save statistics
+        # data_frame = pd.DataFrame(data=results, index=range(1, epoch + 1))
+        # data_frame.to_csv(save_name, index_label='epoch')
+        if test_acc_1 > best_acc:
+           best_acc = test_acc_1
+        wandb.log(
+                {
+                "train_loss": train_loss,
+                "train_acc@1": train_acc_1,
+                "train_acc@5": train_acc_5,
+                "test_loss": test_loss,
+                "test_acc@1": test_acc_1,
+                "test_acc@5": test_acc_5,
+                "best_acc": best_acc
+                }
+            )
+        scheduler.step()
+    wandb.finish()

hubconf.py

@@ -0,0 +1,19 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torchvision.models.resnet import resnet50 as _resnet50
+
+dependencies = ['torch', 'torchvision']
+
+
+def resnet50(pretrained=True, **kwargs):
+    model = _resnet50(pretrained=False, **kwargs)
+    if pretrained:
+        url = 'https://dl.fbaipublicfiles.com/barlowtwins/ep1000_bs2048_lrw0.2_lrb0.0048_lambd0.0051/resnet50.pth'
+        state_dict = torch.hub.load_state_dict_from_url(url, map_location='cpu')
+        model.load_state_dict(state_dict, strict=False)
+    return model

linear.py

@@ -0,0 +1,166 @@
+import argparse
+
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from thop import profile, clever_format
+from torch.utils.data import DataLoader
+from torchvision.datasets import CIFAR10, CIFAR100
+from tqdm import tqdm
+
+import utils
+
+import wandb
+
+import torchvision
+
+class Net(nn.Module):
+    def __init__(self, num_class, pretrained_path, dataset, arch):
+        super(Net, self).__init__()
+
+        if arch=='resnet18':
+            embedding_size = 512
+        elif arch=='resnet50':
+            embedding_size = 2048
+        else:
+            raise NotImplementedError
+
+        # encoder
+        from model import Model
+        self.f = Model(dataset=dataset, arch=arch).f
+        # classifier
+        self.fc = nn.Linear(embedding_size, num_class, bias=True)
+        self.load_state_dict(torch.load(pretrained_path, map_location='cpu'), strict=False)
+
+    def forward(self, x):
+        x = self.f(x)
+        feature = torch.flatten(x, start_dim=1)
+        out = self.fc(feature)
+        return out
+
+# train or test for one epoch
+def train_val(net, data_loader, train_optimizer):
+    is_train = train_optimizer is not None
+    net.train() if is_train else net.eval()
+
+    total_loss, total_correct_1, total_correct_5, total_num, data_bar = 0.0, 0.0, 0.0, 0, tqdm(data_loader)
+    with (torch.enable_grad() if is_train else torch.no_grad()):
+        for data, target in data_bar:
+            data, target = data.cuda(non_blocking=True), target.cuda(non_blocking=True)
+            out = net(data)
+            loss = loss_criterion(out, target)
+
+            if is_train:
+                train_optimizer.zero_grad()
+                loss.backward()
+                train_optimizer.step()
+
+            total_num += data.size(0)
+            total_loss += loss.item() * data.size(0)
+            prediction = torch.argsort(out, dim=-1, descending=True)
+            total_correct_1 += torch.sum((prediction[:, 0:1] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
+            total_correct_5 += torch.sum((prediction[:, 0:5] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
+
+            data_bar.set_description('{} Epoch: [{}/{}] Loss: {:.4f} ACC@1: {:.2f}% ACC@5: {:.2f}% model: {}'
+                                     .format('Train' if is_train else 'Test', epoch, epochs, total_loss / total_num,
+                                             total_correct_1 / total_num * 100, total_correct_5 / total_num * 100,
+                                             model_path.split('/')[-1]))
+
+    return total_loss / total_num, total_correct_1 / total_num * 100, total_correct_5 / total_num * 100
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Linear Evaluation')
+    parser.add_argument('--dataset', default='cifar10', type=str, help='Dataset: cifar10 or tiny_imagenet or stl10')
+    parser.add_argument('--arch', default='resnet50', type=str, help='Backbone architecture for experiments', choices=['resnet50', 'resnet18'])
+    parser.add_argument('--model_path', type=str, default='results/Barlow_Twins/0.005_64_128_model.pth',
+                        help='The base string of the pretrained model path')
+    parser.add_argument('--batch_size', type=int, default=512, help='Number of images in each mini-batch')
+    parser.add_argument('--epochs', type=int, default=200, help='Number of sweeps over the dataset to train')
+
+    args = parser.parse_args()
+
+    wandb.init(project=f"Barlow-Twins-MixUp-Linear-{args.dataset}-{args.arch}", config=args, dir='/data/wbandar1/projects/ssl-aug-artifacts/wandb_logs/')
+    run_id = wandb.run.id
+
+    model_path, batch_size, epochs = args.model_path, args.batch_size, args.epochs
+    dataset = args.dataset
+    if dataset == 'cifar10':
+        train_data = CIFAR10(root='data', train=True,\
+            transform=utils.CifarPairTransform(train_transform = True, pair_transform=False), download=True)
+        test_data = CIFAR10(root='data', train=False,\
+            transform=utils.CifarPairTransform(train_transform = False, pair_transform=False), download=True)
+    if dataset == 'cifar100':
+        train_data = CIFAR100(root='data', train=True,\
+            transform=utils.CifarPairTransform(train_transform = True, pair_transform=False), download=True)
+        test_data = CIFAR100(root='data', train=False,\
+            transform=utils.CifarPairTransform(train_transform = False, pair_transform=False), download=True)
+    elif dataset == 'stl10':
+        train_data =  torchvision.datasets.STL10(root='data', split="train", \
+            transform=utils.StlPairTransform(train_transform = True, pair_transform=False), download=True)
+        test_data =  torchvision.datasets.STL10(root='data', split="test", \
+            transform=utils.StlPairTransform(train_transform = False, pair_transform=False), download=True)
+    elif dataset == 'tiny_imagenet':
+        train_data = torchvision.datasets.ImageFolder('/data/wbandar1/datasets/tiny-imagenet-200/train', \
+                            utils.TinyImageNetPairTransform(train_transform=True, pair_transform=False))
+        test_data = torchvision.datasets.ImageFolder('/data/wbandar1/datasets/tiny-imagenet-200/val', \
+                            utils.TinyImageNetPairTransform(train_transform = False, pair_transform=False))
+
+    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True)
+    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)
+
+    model = Net(num_class=len(train_data.classes), pretrained_path=model_path, dataset=dataset, arch=args.arch).cuda()
+    for param in model.f.parameters():
+        param.requires_grad = False
+
+    if dataset == 'cifar10' or dataset == 'cifar100':
+        flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32).cuda(),))
+    elif dataset == 'tiny_imagenet' or dataset == 'stl10':
+        flops, params = profile(model, inputs=(torch.randn(1, 3, 64, 64).cuda(),))
+    flops, params = clever_format([flops, params])
+    print('# Model Params: {} FLOPs: {}'.format(params, flops))
+
+    # optimizer with lr sheduler
+    lr_start, lr_end = 1e-2, 1e-6
+    gamma = (lr_end / lr_start) ** (1 / epochs)
+    optimizer = optim.Adam(model.fc.parameters(), lr=lr_start, weight_decay=5e-6)
+    scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma)
+    
+    # optimizer with no sheuduler
+    # optimizer = optim.Adam(model.fc.parameters(), lr=1e-3, weight_decay=1e-6)
+
+    loss_criterion = nn.CrossEntropyLoss()
+    results = {'train_loss': [], 'train_acc@1': [], 'train_acc@5': [],
+               'test_loss': [], 'test_acc@1': [], 'test_acc@5': []}
+
+    save_name = model_path.split('.pth')[0] + '_linear.csv'
+
+    best_acc = 0.0
+    for epoch in range(1, epochs + 1):
+        train_loss, train_acc_1, train_acc_5 = train_val(model, train_loader, optimizer)
+        scheduler.step()
+        results['train_loss'].append(train_loss)
+        results['train_acc@1'].append(train_acc_1)
+        results['train_acc@5'].append(train_acc_5)
+        test_loss, test_acc_1, test_acc_5 = train_val(model, test_loader, None)
+        results['test_loss'].append(test_loss)
+        results['test_acc@1'].append(test_acc_1)
+        results['test_acc@5'].append(test_acc_5)
+        # save statistics
+        # data_frame = pd.DataFrame(data=results, index=range(1, epoch + 1))
+        # data_frame.to_csv(save_name, index_label='epoch')
+        #if test_acc_1 > best_acc:
+        #    best_acc = test_acc_1
+        #    torch.save(model.state_dict(), 'results/linear_model.pth')
+        wandb.log(
+                {
+                "train_loss": train_loss,
+                "train_acc@1": train_acc_1,
+                "train_acc@5": train_acc_5,
+                "test_loss": test_loss,
+                "test_acc@1": test_acc_1,
+                "test_acc@5": test_acc_5
+                }
+            )
+    wandb.finish()

main.py

@@ -0,0 +1,271 @@
+import argparse
+import os
+
+import pandas as pd
+import torch
+import numpy as np
+import torch.optim as optim
+import torch.nn.functional as F
+from thop import profile, clever_format
+from torch.utils.data import DataLoader
+from torch.optim.lr_scheduler import MultiStepLR, CosineAnnealingWarmRestarts
+from tqdm import tqdm
+
+import utils
+from model import Model
+import math
+
+import torchvision
+
+import wandb
+
+if torch.cuda.is_available():
+    torch.backends.cudnn.benchmark = True
+
+def off_diagonal(x):
+    n, m = x.shape
+    assert n == m
+    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
+
+def adjust_learning_rate(args, optimizer, loader, step):
+    max_steps = args.epochs * len(loader)
+    warmup_steps = 10 * len(loader)
+    base_lr = args.batch_size / 256
+    if step < warmup_steps:
+        lr = base_lr * step / warmup_steps
+    else:
+        step -= warmup_steps
+        max_steps -= warmup_steps
+        q = 0.5 * (1 + math.cos(math.pi * step / max_steps))
+        end_lr = base_lr * 0.001
+        lr = base_lr * q + end_lr * (1 - q)
+    optimizer.param_groups[0]['lr'] = lr * args.lr
+
+def train(args, epoch, net, data_loader, train_optimizer):
+    net.train()
+    total_loss, total_loss_bt, total_loss_mix, total_num, train_bar = 0.0, 0.0, 0.0, 0, tqdm(data_loader)
+    for step, data_tuple in enumerate(train_bar, start=epoch * len(train_bar)):
+        if args.lr_shed == "cosine":
+            adjust_learning_rate(args, train_optimizer, data_loader, step)
+        (pos_1, pos_2), _ = data_tuple
+        pos_1, pos_2 = pos_1.cuda(non_blocking=True), pos_2.cuda(non_blocking=True)
+        _, out_1 = net(pos_1)
+        _, out_2 = net(pos_2)
+        
+        out_1_norm = (out_1 - out_1.mean(dim=0)) / out_1.std(dim=0)
+        out_2_norm = (out_2 - out_2.mean(dim=0)) / out_2.std(dim=0)
+        c = torch.matmul(out_1_norm.T, out_2_norm) / batch_size
+        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
+        off_diag = off_diagonal(c).pow_(2).sum()
+        loss_bt = on_diag + lmbda * off_diag
+        
+        ##  MixUp (Our Contribution) ## 
+        if args.is_mixup.lower() == 'true':
+            index = torch.randperm(batch_size).cuda(non_blocking=True)
+            alpha = np.random.beta(1.0, 1.0)
+            pos_m = alpha * pos_1 + (1 - alpha) * pos_2[index, :]
+            
+            _, out_m = net(pos_m)
+            out_m_norm = (out_m - out_m.mean(dim=0)) / out_m.std(dim=0)
+            
+            cc_m_1 = torch.matmul(out_m_norm.T, out_1_norm) / batch_size
+            cc_m_1_gt = alpha*torch.matmul(out_1_norm.T, out_1_norm) / batch_size + \
+                            (1-alpha)*torch.matmul(out_2_norm[index,:].T, out_1_norm) / batch_size
+
+            cc_m_2 = torch.matmul(out_m_norm.T, out_2_norm) / batch_size
+            cc_m_2_gt = alpha*torch.matmul(out_1_norm.T, out_2_norm) / batch_size + \
+                            (1-alpha)*torch.matmul(out_2_norm[index,:].T, out_2_norm) / batch_size
+            
+            loss_mix = args.mixup_loss_scale*lmbda*((cc_m_1-cc_m_1_gt).pow_(2).sum() + (cc_m_2-cc_m_2_gt).pow_(2).sum())
+        else:
+            loss_mix = torch.zeros(1).cuda()
+        ##  MixUp (Our Contribution) ##
+        
+        loss = loss_bt + loss_mix
+        train_optimizer.zero_grad()
+        loss.backward()
+        train_optimizer.step()
+
+        total_num += batch_size
+        total_loss += loss.item() * batch_size
+        total_loss_bt += loss_bt.item() * batch_size
+        total_loss_mix += loss_mix.item() * batch_size
+        
+        train_bar.set_description('Train Epoch: [{}/{}] lr: {:.3f}x10-3 Loss: {:.4f} lmbda:{:.4f} bsz:{} f_dim:{} dataset: {}'.format(\
+                                epoch, epochs, train_optimizer.param_groups[0]['lr'] * 1000, total_loss / total_num, lmbda, batch_size, feature_dim, dataset))
+    return total_loss_bt / total_num, total_loss_mix / total_num, total_loss / total_num
+
+
+def test(net, memory_data_loader, test_data_loader):
+    net.eval()
+    total_top1, total_top5, total_num, feature_bank, target_bank = 0.0, 0.0, 0, [], []
+    with torch.no_grad():
+        # generate feature bank and target bank
+        for data_tuple in tqdm(memory_data_loader, desc='Feature extracting'):
+            (data, _), target = data_tuple
+            target_bank.append(target)
+            feature, out = net(data.cuda(non_blocking=True))
+            feature_bank.append(feature)
+        # [D, N]
+        feature_bank = torch.cat(feature_bank, dim=0).t().contiguous()
+        # [N]
+        feature_labels = torch.cat(target_bank, dim=0).contiguous().to(feature_bank.device)
+        # loop test data to predict the label by weighted knn search
+        test_bar = tqdm(test_data_loader)
+        for data_tuple in test_bar:
+            (data, _), target = data_tuple
+            data, target = data.cuda(non_blocking=True), target.cuda(non_blocking=True)
+            feature, out = net(data)
+
+            total_num += data.size(0)
+            # compute cos similarity between each feature vector and feature bank ---> [B, N]
+            sim_matrix = torch.mm(feature, feature_bank)
+            # [B, K]
+            sim_weight, sim_indices = sim_matrix.topk(k=k, dim=-1)
+            # [B, K]
+            sim_labels = torch.gather(feature_labels.expand(data.size(0), -1), dim=-1, index=sim_indices)
+            sim_weight = (sim_weight / temperature).exp()
+
+            # counts for each class
+            one_hot_label = torch.zeros(data.size(0) * k, c, device=sim_labels.device)
+            # [B*K, C]
+            one_hot_label = one_hot_label.scatter(dim=-1, index=sim_labels.view(-1, 1), value=1.0)
+            # weighted score ---> [B, C]
+            pred_scores = torch.sum(one_hot_label.view(data.size(0), -1, c) * sim_weight.unsqueeze(dim=-1), dim=1)
+
+            pred_labels = pred_scores.argsort(dim=-1, descending=True)
+            total_top1 += torch.sum((pred_labels[:, :1] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
+            total_top5 += torch.sum((pred_labels[:, :5] == target.unsqueeze(dim=-1)).any(dim=-1).float()).item()
+            test_bar.set_description('Test Epoch: [{}/{}] Acc@1:{:.2f}% Acc@5:{:.2f}%'
+                                     .format(epoch, epochs, total_top1 / total_num * 100, total_top5 / total_num * 100))
+    return total_top1 / total_num * 100, total_top5 / total_num * 100
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Training Barlow Twins')
+    parser.add_argument('--dataset', default='cifar10', type=str, help='Dataset: cifar10, cifar100, tiny_imagenet, stl10', choices=['cifar10', 'cifar100', 'tiny_imagenet', 'stl10'])
+    parser.add_argument('--arch', default='resnet50', type=str, help='Backbone architecture', choices=['resnet50', 'resnet18'])
+    parser.add_argument('--feature_dim', default=128, type=int, help='Feature dim for embedding vector')
+    parser.add_argument('--temperature', default=0.5, type=float, help='Temperature used in softmax (kNN evaluation)')
+    parser.add_argument('--k', default=200, type=int, help='Top k most similar images used to predict the label')
+    parser.add_argument('--batch_size', default=512, type=int, help='Number of images in each mini-batch')
+    parser.add_argument('--epochs', default=1000, type=int, help='Number of sweeps over the dataset to train')
+    parser.add_argument('--lr', default=1e-3, type=float, help='Base learning rate')
+    parser.add_argument('--lr_shed', default="step", choices=["step", "cosine"], type=str, help='Learning rate scheduler: step / cosine')
+    
+    # for barlow twins
+    parser.add_argument('--lmbda', default=0.005, type=float, help='Lambda that controls the on- and off-diagonal terms')
+    parser.add_argument('--corr_neg_one', dest='corr_neg_one', action='store_true')
+    parser.add_argument('--corr_zero', dest='corr_neg_one', action='store_false')
+    parser.set_defaults(corr_neg_one=False)
+    
+    # for mixup
+    parser.add_argument('--is_mixup', dest='is_mixup', type=str, default='false', choices=['true', 'false'])
+    parser.add_argument('--mixup_loss_scale', dest='mixup_loss_scale', type=float, default=5.0)
+
+    # GPU id (just for record)
+    parser.add_argument('--gpu', dest='gpu', type=int, default=0)
+    
+    args = parser.parse_args()
+    is_mixup = args.is_mixup.lower() == 'true'
+
+    wandb.init(project=f"Barlow-Twins-MixUp-{args.dataset}-{args.arch}", config=args, dir='results/wandb_logs/')
+    run_id = wandb.run.id
+    dataset = args.dataset
+    feature_dim, temperature, k = args.feature_dim, args.temperature, args.k
+    batch_size, epochs = args.batch_size, args.epochs
+    lmbda = args.lmbda
+    corr_neg_one = args.corr_neg_one
+    
+    if dataset == 'cifar10':
+        train_data = torchvision.datasets.CIFAR10(root='/data/wbandar1/datasets', train=True, \
+                                                  transform=utils.CifarPairTransform(train_transform = True), download=True)
+        memory_data = torchvision.datasets.CIFAR10(root='/data/wbandar1/datasets', train=True, \
+                                                  transform=utils.CifarPairTransform(train_transform = False), download=True)
+        test_data = torchvision.datasets.CIFAR10(root='/data/wbandar1/datasets', train=False, \
+                                                  transform=utils.CifarPairTransform(train_transform = False), download=True)
+    elif dataset == 'cifar100':
+        train_data = torchvision.datasets.CIFAR100(root='/data/wbandar1/datasets', train=True, \
+                                                  transform=utils.CifarPairTransform(train_transform = True), download=True)
+        memory_data = torchvision.datasets.CIFAR100(root='/data/wbandar1/datasets', train=True, \
+                                                    transform=utils.CifarPairTransform(train_transform = False), download=True)
+        test_data = torchvision.datasets.CIFAR100(root='/data/wbandar1/datasets', train=False, \
+                                                    transform=utils.CifarPairTransform(train_transform = False), download=True)
+    elif dataset == 'stl10':
+        train_data = torchvision.datasets.STL10(root='/data/wbandar1/datasets', split="train+unlabeled", \
+                                                    transform=utils.StlPairTransform(train_transform = True), download=True)
+        memory_data = torchvision.datasets.STL10(root='/data/wbandar1/datasets', split="train", \
+                                                  transform=utils.StlPairTransform(train_transform = False), download=True)
+        test_data = torchvision.datasets.STL10(root='/data/wbandar1/datasets', split="test", \
+                                                  transform=utils.StlPairTransform(train_transform = False), download=True)
+    elif dataset == 'tiny_imagenet':
+        # download if not exits
+        if not os.path.isdir('/data/wbandar1/datasets/tiny-imagenet-200'):
+            raise ValueError("First preprocess the tinyimagenet dataset...") 
+            
+        train_data = torchvision.datasets.ImageFolder('/data/wbandar1/datasets/tiny-imagenet-200/train', \
+                                                        utils.TinyImageNetPairTransform(train_transform = True))
+        memory_data = torchvision.datasets.ImageFolder('/data/wbandar1/datasets/tiny-imagenet-200/train', \
+                                                      utils.TinyImageNetPairTransform(train_transform = False))
+        test_data = torchvision.datasets.ImageFolder('/data/wbandar1/datasets/tiny-imagenet-200/val', \
+                                                      utils.TinyImageNetPairTransform(train_transform = False))
+    
+    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True,
+                            drop_last=True)
+    memory_loader = DataLoader(memory_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)
+    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)
+
+    # model setup and optimizer config
+    model = Model(feature_dim, dataset, args.arch).cuda()
+    if dataset == 'cifar10' or dataset == 'cifar100':
+        flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32).cuda(),))
+    elif dataset == 'tiny_imagenet' or dataset == 'stl10':
+        flops, params = profile(model, inputs=(torch.randn(1, 3, 64, 64).cuda(),))
+    flops, params = clever_format([flops, params])
+    print('# Model Params: {} FLOPs: {}'.format(params, flops))
+    
+    optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-6)
+    if args.lr_shed == "step":
+        m = [args.epochs - a for a in [50, 25]]
+        scheduler = MultiStepLR(optimizer, milestones=m, gamma=0.2)
+    c = len(memory_data.classes)
+
+    results = {'train_loss': [], 'test_acc@1': [], 'test_acc@5': []}
+    save_name_pre = '{}_{}_{}_{}_{}'.format(run_id, lmbda, feature_dim, batch_size, dataset)
+    run_id_dir = os.path.join('results/', run_id)
+    if not os.path.exists(run_id_dir):
+        print('Creating directory {}'.format(run_id_dir))
+        os.mkdir(run_id_dir)
+
+    best_acc = 0.0
+    for epoch in range(1, epochs + 1):
+        loss_bt, loss_mix, train_loss = train(args, epoch, model, train_loader, optimizer)
+        if args.lr_shed == "step":
+            scheduler.step()
+        wandb.log(
+                {
+                "epoch": epoch,
+                "lr": optimizer.param_groups[0]['lr'],
+                "loss_bt": loss_bt,
+                "loss_mix": loss_mix,
+                "train_loss": train_loss}
+            )
+        if epoch % 5 == 0:
+            test_acc_1, test_acc_5 = test(model, memory_loader, test_loader)
+
+            results['train_loss'].append(train_loss)
+            results['test_acc@1'].append(test_acc_1)
+            results['test_acc@5'].append(test_acc_5)
+            data_frame = pd.DataFrame(data=results, index=range(5, epoch + 1, 5))
+            data_frame.to_csv('results/{}_statistics.csv'.format(save_name_pre), index_label='epoch')
+            wandb.log(
+                {
+                "test_acc@1": test_acc_1,
+                "test_acc@5": test_acc_5
+                }
+            )
+            if test_acc_1 > best_acc:
+                best_acc = test_acc_1
+                torch.save(model.state_dict(), 'results/{}/{}_model.pth'.format(run_id, save_name_pre))
+        if epoch % 50 == 0:
+            torch.save(model.state_dict(), 'results/{}/{}_model_{}.pth'.format(run_id, save_name_pre, epoch))
+    wandb.finish()

main_imagenet.py

@@ -0,0 +1,463 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from pathlib import Path
+import argparse
+import json
+import math
+import os
+import random
+import signal
+import subprocess
+import sys
+import time
+import numpy as np
+import wandb
+
+from PIL import Image, ImageOps, ImageFilter
+from torch import nn, optim
+import torch
+import torchvision
+import torchvision.transforms as transforms
+
+parser = argparse.ArgumentParser(description='Barlow Twins Training')
+parser.add_argument('data', type=Path, metavar='DIR',
+                    help='path to dataset')
+parser.add_argument('--workers', default=8, type=int, metavar='N',
+                    help='number of data loader workers')
+parser.add_argument('--epochs', default=300, type=int, metavar='N',
+                    help='number of total epochs to run')
+parser.add_argument('--batch-size', default=512, type=int, metavar='N',
+                    help='mini-batch size')
+parser.add_argument('--learning-rate-weights', default=0.2, type=float, metavar='LR',
+                    help='base learning rate for weights')
+parser.add_argument('--learning-rate-biases', default=0.0048, type=float, metavar='LR',
+                    help='base learning rate for biases and batch norm parameters')
+parser.add_argument('--weight-decay', default=1e-6, type=float, metavar='W',
+                    help='weight decay')
+parser.add_argument('--lambd', default=0.0051, type=float, metavar='L',
+                    help='weight on off-diagonal terms')
+parser.add_argument('--projector', default='8192-8192-8192', type=str,
+                    metavar='MLP', help='projector MLP')
+parser.add_argument('--print-freq', default=1, type=int, metavar='N',
+                    help='print frequency')
+parser.add_argument('--checkpoint-dir', default='/mnt/store/wbandar1/projects/ssl-aug-artifacts/', type=Path,
+                    metavar='DIR', help='path to checkpoint directory')
+parser.add_argument('--is_mixup', default='false', type=str,
+                    metavar='L', help='mixup regularization', choices=['true', 'false'])
+parser.add_argument('--lambda_mixup', default=0.1, type=float, metavar='L',
+                    help='Hyperparamter for the regularization loss')
+
+def main():
+    args = parser.parse_args()
+    args.is_mixup = args.is_mixup.lower() == 'true'
+    args.ngpus_per_node = torch.cuda.device_count()
+
+    run = wandb.init(project="Barlow-Twins-MixUp-ImageNet", config=args, dir='/mnt/store/wbandar1/projects/ssl-aug-artifacts/wandb_logs/')
+    run_id = wandb.run.id
+    args.checkpoint_dir=Path(os.path.join(args.checkpoint_dir, run_id))
+
+    if 'SLURM_JOB_ID' in os.environ:
+        # single-node and multi-node distributed training on SLURM cluster
+        # requeue job on SLURM preemption
+        signal.signal(signal.SIGUSR1, handle_sigusr1)
+        signal.signal(signal.SIGTERM, handle_sigterm)
+        # find a common host name on all nodes
+        # assume scontrol returns hosts in the same order on all nodes
+        cmd = 'scontrol show hostnames ' + os.getenv('SLURM_JOB_NODELIST')
+        stdout = subprocess.check_output(cmd.split())
+        host_name = stdout.decode().splitlines()[0]
+        args.rank = int(os.getenv('SLURM_NODEID')) * args.ngpus_per_node
+        args.world_size = int(os.getenv('SLURM_NNODES')) * args.ngpus_per_node
+        args.dist_url = f'tcp://{host_name}:58472'
+    else:
+        # single-node distributed training
+        args.rank = 0
+        args.dist_url = 'tcp://localhost:58472'
+        args.world_size = args.ngpus_per_node
+    torch.multiprocessing.spawn(main_worker, (args,run,), args.ngpus_per_node)
+    wandb.finish()
+
+
+def main_worker(gpu, args, run):
+    args.rank += gpu
+    torch.distributed.init_process_group(
+        backend='nccl', init_method=args.dist_url,
+        world_size=args.world_size, rank=args.rank)
+
+    if args.rank == 0:
+        args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        stats_file = open(args.checkpoint_dir / 'stats.txt', 'a', buffering=1)
+        print(' '.join(sys.argv))
+        print(' '.join(sys.argv), file=stats_file)
+
+    torch.cuda.set_device(gpu)
+    torch.backends.cudnn.benchmark = True
+
+    model = BarlowTwins(args).cuda(gpu)
+    model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
+    param_weights = []
+    param_biases = []
+    for param in model.parameters():
+        if param.ndim == 1:
+            param_biases.append(param)
+        else:
+            param_weights.append(param)
+    parameters = [{'params': param_weights}, {'params': param_biases}]
+    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
+    optimizer = LARS(parameters, lr=0, weight_decay=args.weight_decay,
+                     weight_decay_filter=True,
+                     lars_adaptation_filter=True)
+
+    # automatically resume from checkpoint if it exists
+    if (args.checkpoint_dir / 'checkpoint.pth').is_file():
+        ckpt = torch.load(args.checkpoint_dir / 'checkpoint.pth',
+                          map_location='cpu')
+        start_epoch = ckpt['epoch']
+        model.load_state_dict(ckpt['model'])
+        optimizer.load_state_dict(ckpt['optimizer'])
+    else:
+        start_epoch = 0
+
+    dataset = torchvision.datasets.ImageFolder(args.data / 'train', Transform())
+    sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+    assert args.batch_size % args.world_size == 0
+    per_device_batch_size = args.batch_size // args.world_size
+    loader = torch.utils.data.DataLoader(
+        dataset, batch_size=per_device_batch_size, num_workers=args.workers,
+        pin_memory=True, sampler=sampler)
+
+    start_time = time.time()
+    scaler = torch.cuda.amp.GradScaler(growth_interval=100, enabled=True)
+    for epoch in range(start_epoch, args.epochs):
+        sampler.set_epoch(epoch)
+        for step, ((y1, y2), _) in enumerate(loader, start=epoch * len(loader)):
+            y1 = y1.cuda(gpu, non_blocking=True)
+            y2 = y2.cuda(gpu, non_blocking=True)
+            adjust_learning_rate(args, optimizer, loader, step)
+            mixup_loss_scale = adjust_mixup_scale(loader, step, args.lambda_mixup)
+            optimizer.zero_grad()
+            with torch.cuda.amp.autocast(enabled=True):
+                loss_bt, loss_reg = model(y1, y2, args.is_mixup)
+            loss_regs = mixup_loss_scale * loss_reg
+            loss = loss_bt + loss_regs
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            if step % args.print_freq == 0:
+                if args.rank == 0:
+                    stats = dict(epoch=epoch, step=step,
+                                 lr_weights=optimizer.param_groups[0]['lr'],
+                                 lr_biases=optimizer.param_groups[1]['lr'],
+                                 loss=loss.item(),
+                                 time=int(time.time() - start_time))
+                    print(json.dumps(stats))
+                    print(json.dumps(stats), file=stats_file)
+                    if args.is_mixup:
+                        run.log(
+                            {
+                                "epoch": epoch,
+                                "step": step,
+                                "lr_weights": optimizer.param_groups[0]['lr'],
+                                "lr_biases": optimizer.param_groups[1]['lr'],
+                                "loss": loss.item(),
+                                "loss_bt": loss_bt.item(),
+                                "loss_reg(unscaled)": loss_reg.item(),
+                                "reg_scale": mixup_loss_scale,
+                                "loss_reg(scaled)": loss_regs.item(),
+                                "time": int(time.time() - start_time)}
+                                )
+                    else:
+                        run.log(
+                            {
+                                "epoch": epoch,
+                                "step": step,
+                                "lr_weights": optimizer.param_groups[0]['lr'],
+                                "lr_biases": optimizer.param_groups[1]['lr'],
+                                "loss": loss.item(),
+                                "loss_bt": loss.item(),
+                                "loss_reg(unscaled)": 0.,
+                                "reg_scale": 0.,
+                                "loss_reg(scaled)": 0.,
+                                "time": int(time.time() - start_time)}
+                                )
+        if args.rank == 0:
+            # save checkpoint
+            state = dict(epoch=epoch + 1, model=model.state_dict(),
+                         optimizer=optimizer.state_dict())
+            torch.save(state, args.checkpoint_dir / 'checkpoint.pth')
+    if args.rank == 0:
+        # save final model
+        print("Saving final model ...")
+        torch.save(model.module.backbone.state_dict(),
+                   args.checkpoint_dir / 'resnet50.pth')
+        print("Finished saving final model ...")
+
+
+def adjust_learning_rate(args, optimizer, loader, step):
+    max_steps = args.epochs * len(loader)
+    warmup_steps = 10 * len(loader)
+    base_lr = args.batch_size / 256
+    if step < warmup_steps:
+        lr = base_lr * step / warmup_steps
+    else:
+        step -= warmup_steps
+        max_steps -= warmup_steps
+        q = 0.5 * (1 + math.cos(math.pi * step / max_steps))
+        end_lr = base_lr * 0.001
+        lr = base_lr * q + end_lr * (1 - q)
+    optimizer.param_groups[0]['lr'] = lr * args.learning_rate_weights
+    optimizer.param_groups[1]['lr'] = lr * args.learning_rate_biases
+
+def adjust_mixup_scale(loader, step, lambda_mixup):
+    warmup_steps = 10 * len(loader)
+    if step < warmup_steps:
+        return lambda_mixup * step / warmup_steps
+    else:
+        return lambda_mixup
+
+def handle_sigusr1(signum, frame):
+    os.system(f'scontrol requeue {os.getenv("SLURM_JOB_ID")}')
+    exit()
+
+
+def handle_sigterm(signum, frame):
+    pass
+
+
+def off_diagonal(x):
+    # return a flattened view of the off-diagonal elements of a square matrix
+    n, m = x.shape
+    assert n == m
+    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
+
+
+class BarlowTwins(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.backbone = torchvision.models.resnet50(zero_init_residual=True)
+        self.backbone.fc = nn.Identity()
+
+        # projector
+        sizes = [2048] + list(map(int, args.projector.split('-')))
+        layers = []
+        for i in range(len(sizes) - 2):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=False))
+            layers.append(nn.BatchNorm1d(sizes[i + 1]))
+            layers.append(nn.ReLU(inplace=True))
+        layers.append(nn.Linear(sizes[-2], sizes[-1], bias=False))
+        self.projector = nn.Sequential(*layers)
+
+        # normalization layer for the representations z1 and z2
+        # self.bn = nn.BatchNorm1d(sizes[-1], affine=False)
+
+    # def forward(self, y1, y2):
+    #     z1 = self.projector(self.backbone(y1))
+    #     z2 = self.projector(self.backbone(y2))
+
+    #     # empirical cross-correlation matrix
+    #     c = self.bn(z1).T @ self.bn(z2)
+
+    #     # sum the cross-correlation matrix between all gpus
+    #     c.div_(self.args.batch_size)
+    #     torch.distributed.all_reduce(c)
+
+    #     on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
+    #     off_diag = off_diagonal(c).pow_(2).sum()
+    #     loss = on_diag + self.args.lambd * off_diag
+    #     return loss
+
+    def forward(self, y1, y2, is_mixup):
+        batch_size = y1.shape[0]
+
+        ### original barlow twins ###
+        z1 = self.projector(self.backbone(y1))
+        z2 = self.projector(self.backbone(y2))
+
+        # normilization
+        z1 = (z1 - z1.mean(dim=0)) / z1.std(dim=0)
+        z2 = (z2 - z2.mean(dim=0)) / z2.std(dim=0)
+
+        # empirical cross-correlation matrix
+        c = z1.T @ z2
+
+        # sum the cross-correlation matrix between all gpus
+        c.div_(self.args.batch_size)
+        torch.distributed.all_reduce(c)
+
+        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
+        off_diag = off_diagonal(c).pow_(2).sum()
+        loss = on_diag + self.args.lambd * off_diag
+
+        if is_mixup:
+            ##############################################
+            ### mixup regularization: Implementation 1 ###
+            ##############################################
+            
+            # index = torch.randperm(batch_size).cuda(non_blocking=True)
+            # alpha = np.random.beta(1.0, 1.0)
+            # ym = alpha * y1 + (1. - alpha) * y2[index, :]
+            # zm = self.projector(self.backbone(ym))
+
+            # # normilization
+            # zm = (zm - zm.mean(dim=0)) / zm.std(dim=0)
+
+            # # cc
+            # cc_m_1 = zm.T @ z1
+            # cc_m_1.div_(batch_size)
+            # cc_m_1_gt = alpha*(z1.T @ z1) + (1.-alpha)*(z2[index,:].T @ z1)
+            # cc_m_1_gt.div_(batch_size)
+
+            # cc_m_2 = zm.T @ z2
+            # cc_m_2.div_(batch_size)
+            # cc_m_2_gt = alpha*(z2.T @ z2) + (1.-alpha)*(z2[index,:].T @ z2)
+            # cc_m_2_gt.div_(batch_size)
+
+            # # mixup reg. loss
+            # lossm = 0.5*self.args.lambd*((cc_m_1-cc_m_1_gt).pow_(2).sum() + (cc_m_2-cc_m_2_gt).pow_(2).sum())
+            
+            ##############################################
+            ### mixup regularization: Implementation 2 ###
+            ##############################################
+            index = torch.randperm(batch_size).cuda(non_blocking=True)
+            alpha = np.random.beta(1.0, 1.0)
+            ym = alpha * y1 + (1. - alpha) * y2[index, :]
+            zm = self.projector(self.backbone(ym))
+
+            # normilization
+            zm = (zm - zm.mean(dim=0)) / zm.std(dim=0)
+
+            # cc
+            cc_m_1 = zm.T @ z1
+            cc_m_1.div_(self.args.batch_size)
+            cc_m_1_gt = alpha*(z1.T @ z1) + (1.-alpha)*(z2[index,:].T @ z1)
+            cc_m_1_gt.div_(self.args.batch_size)
+
+            cc_m_2 = zm.T @ z2
+            cc_m_2.div_(self.args.batch_size)
+            cc_m_2_gt = alpha*(z2.T @ z2) + (1.-alpha)*(z2[index,:].T @ z2)
+            cc_m_2_gt.div_(self.args.batch_size)
+
+            # gathering all cc
+            torch.distributed.all_reduce(cc_m_1)
+            torch.distributed.all_reduce(cc_m_1_gt)
+            torch.distributed.all_reduce(cc_m_2)
+            torch.distributed.all_reduce(cc_m_2_gt)
+
+            # mixup reg. loss
+            lossm = 0.5*self.args.lambd*((cc_m_1-cc_m_1_gt).pow_(2).sum() + (cc_m_2-cc_m_2_gt).pow_(2).sum())
+        else:
+            lossm = torch.zeros(1)
+        return loss, lossm
+
+class LARS(optim.Optimizer):
+    def __init__(self, params, lr, weight_decay=0, momentum=0.9, eta=0.001,
+                 weight_decay_filter=False, lars_adaptation_filter=False):
+        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)
+
+
+    def exclude_bias_and_norm(self, p):
+        return p.ndim == 1
+
+    @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 not g['weight_decay_filter'] or not self.exclude_bias_and_norm(p):
+                    dp = dp.add(p, alpha=g['weight_decay'])
+
+                if not g['lars_adaptation_filter'] or not self.exclude_bias_and_norm(p):
+                    param_norm = torch.norm(p)
+                    update_norm = torch.norm(dp)
+                    one = torch.ones_like(param_norm)
+                    q = torch.where(param_norm > 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'])
+
+
+
+class GaussianBlur(object):
+    def __init__(self, p):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            sigma = random.random() * 1.9 + 0.1
+            return img.filter(ImageFilter.GaussianBlur(sigma))
+        else:
+            return img
+
+
+class Solarization(object):
+    def __init__(self, p):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return ImageOps.solarize(img)
+        else:
+            return img
+
+
+class Transform:
+    def __init__(self):
+        self.transform = transforms.Compose([
+            transforms.RandomResizedCrop(224, interpolation=Image.BICUBIC),
+            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),
+            GaussianBlur(p=1.0),
+            Solarization(p=0.0),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+        ])
+        self.transform_prime = transforms.Compose([
+            transforms.RandomResizedCrop(224, interpolation=Image.BICUBIC),
+            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),
+            GaussianBlur(p=0.1),
+            Solarization(p=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+        ])
+
+    def __call__(self, x):
+        y1 = self.transform(x)
+        y2 = self.transform_prime(x)
+        return y1, y2
+
+
+if __name__ == '__main__':
+    main()

model.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models.resnet import resnet50, resnet18
+
+
+class Model(nn.Module):
+    def __init__(self, feature_dim=128, dataset='cifar10', arch='resnet50'):
+        super(Model, self).__init__()
+
+        self.f = []
+        if arch == 'resnet18':
+            temp_model = resnet18().named_children()
+            embedding_size = 512
+        elif arch == 'resnet50':
+            temp_model = resnet50().named_children()
+            embedding_size = 2048
+        else:
+            raise NotImplementedError
+        
+        for name, module in temp_model:
+            if name == 'conv1':
+                module = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
+            if dataset == 'cifar10' or dataset == 'cifar100':
+                if not isinstance(module, nn.Linear) and not isinstance(module, nn.MaxPool2d):
+                    self.f.append(module)
+            elif dataset == 'tiny_imagenet' or dataset == 'stl10':
+                if not isinstance(module, nn.Linear):
+                    self.f.append(module)
+        # encoder
+        self.f = nn.Sequential(*self.f)
+        # projection head
+        self.g = nn.Sequential(nn.Linear(embedding_size, 512, bias=False), nn.BatchNorm1d(512),
+                               nn.ReLU(inplace=True), nn.Linear(512, feature_dim, bias=True))
+
+    def forward(self, x):
+        x = self.f(x)
+        feature = torch.flatten(x, start_dim=1)
+        out = self.g(feature)
+        return F.normalize(feature, dim=-1), F.normalize(out, dim=-1)

preprocess_datasets/preprocess_tinyimagenet.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# download and unzip dataset
+cd /data/wbandar1/datasets
+wget http://cs231n.stanford.edu/tiny-imagenet-200.zip
+unzip tiny-imagenet-200.zip
+
+current="$(pwd)/tiny-imagenet-200"
+
+# training data
+cd $current/train
+for DIR in $(ls); do
+   cd $DIR
+   rm *.txt
+   mv images/* .
+   rm -r images
+   cd ..
+done
+
+# validation data
+cd $current/val
+annotate_file="val_annotations.txt"
+length=$(cat $annotate_file | wc -l)
+for i in $(seq 1 $length); do
+    # fetch i th line
+    line=$(sed -n ${i}p $annotate_file)
+    # get file name and directory name
+    file=$(echo $line | cut -f1 -d" " )
+    directory=$(echo $line | cut -f2 -d" ")
+    mkdir -p $directory
+    mv images/$file $directory
+done
+rm -r images
+echo "done"

scripts-linear-resnet18/cifar10.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=cifar10
+arch=resnet18
+batch_size=512
+model_path=checkpoints/4wdhbpcf_0.0078125_1024_256_cifar10_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet18/cifar100.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=cifar100
+arch=resnet18
+batch_size=512
+model_path=checkpoints/76kk7scz_0.0078125_1024_256_cifar100_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-sug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet18/stl10.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=stl10
+arch=resnet18
+batch_size=512
+model_path=checkpoints/i7det4xq_0.0078125_1024_256_stl10_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-sug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet18/tinyimagenet.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=tiny_imagenet
+arch=resnet18
+batch_size=512
+model_path=checkpoints/02azq6fs_0.0009765_1024_256_tiny_imagenet_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-sug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet50/cifar10.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=cifar10
+arch=resnet50
+batch_size=512
+model_path=checkpoints/v3gwgusq_0.0078125_1024_256_cifar10_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet50/cifar100.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=cifar100
+arch=resnet50
+batch_size=512
+model_path=checkpoints/z6ngefw7_0.0078125_1024_256_cifar100_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet50/imagenet_sup.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+path_to_imagenet_data=datasets/imagenet1k/
+path_to_model=checkpoints/13awtq23_0.0051_8192_1024_imagenet_0.1_resnet50.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "NCCL_P2P_DISABLE=1 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python evaluate_imagenet.py ${path_to_imagenet_data} ${path_to_model} --lr-classifier 0.3^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet50/stl10.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=stl10
+arch=resnet50
+batch_size=512
+model_path=checkpoints/pbknx38b_0.0078125_1024_256_stl10_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-linear-resnet50/tinyimagenet.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+gpu=0
+dataset=tiny_imagenet
+arch=resnet50
+batch_size=512
+model_path=checkpoints/kxlkigsv_0.0009765_1024_256_tiny_imagenet_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python linear.py --dataset ${dataset} --model_path ${model_path} --arch ${arch}^M"
+screen -S "$session_name" -X detachs

scripts-pretrain-resnet18/cifar10.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+# default: https://wandb.ai/cha-yas/Barlow-Twins-MixUp-cifar10-resnet18/runs/4wdhbpcf/overview?workspace=user-wgcban
+gpu=0
+dataset=cifar10
+arch=resnet18
+feature_dim=1024
+is_mixup=true # true, false
+batch_size=256
+epochs=2000
+lr=0.01
+lr_shed=cosine # step, cosine
+mixup_loss_scale=4.0 # scale w.r.t. lambda: 0.0078125 * 5 = 0.0390625
+lmbda=0.0078125
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-pretrain-resnet18/cifar100.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+gpu=0
+dataset=cifar100
+arch=resnet18
+feature_dim=2048
+is_mixup=true # true, false
+batch_size=256
+epochs=2000
+lr=0.01
+lr_shed=cosine #"step", "cosine" # step, cosine
+mixup_loss_scale=4.0 # scale w.r.t. lambda: 0.0078125 * 5 = 0.0390625
+lmbda=0.0078125
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-pretrain-resnet18/stl10.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+gpu=0
+dataset=stl10
+arch=resnet18
+feature_dim=1024
+is_mixup=true # true, false
+batch_size=256
+epochs=2000
+lr=0.01
+lr_shed=cosine #"step", "cosine" # step, cosine
+mixup_loss_scale=2.0 # scale w.r.t. lambda: 0.0078125 * 5 = 0.0390625pochs=2000
+lmbda=0.0078125
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-pretrain-resnet18/tinyimagenet.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+gpu=0
+dataset=tiny_imagenet
+arch=resnet18
+feature_dim=1024
+is_mixup=true # true, false
+batch_size=256
+epochs=2000
+lr=0.01
+lr_shed=cosine #"step", "cosine" # step, cosine
+mixup_loss_scale=4.0 # scale w.r.t. lambda
+lmbda=$(echo "scale=7; 1 / ${feature_dim}" | bc)
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-pretrain-resnet50/cifar10.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+gpu=0
+dataset=cifar10
+arch=resnet50
+feature_dim=1024
+is_mixup=true # true, false
+batch_size=256
+epochs=1000
+lr=0.01
+lr_shed=cosine # step, cosine
+mixup_loss_scale=4.0 # scale w.r.t. lambda: 0.0078125 * 5 = 0.0390625
+lmbda=0.0078125
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-pretrain-resnet50/cifar100.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+gpu=0
+dataset=cifar100
+arch=resnet50
+feature_dim=1024
+is_mixup=true # true, false
+batch_size=256
+epochs=1000
+lr=0.01
+lr_shed=cosine # step, cosine
+mixup_loss_scale=4.0 # scale w.r.t. lambda: 0.0078125 * 5 = 0.0390625
+lmbda=0.0078125
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-pretrain-resnet50/imagenet.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+is_mixup=true
+batch_size=1024 #128/gpu works
+lr_w=0.2 #0.2
+lr_b=0.0048 #0.0048
+lambda_mixup=1.0
+
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "NCCL_P2P_DISABLE=1 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python main_imagenet.py /data/wbandar1/datasets/imagenet1k/ --is_mixup ${is_mixup} --batch-size ${batch_size} --learning-rate-weights ${lr_w} --learning-rate-biases ${lr_b} --lambda_mixup ${lambda_mixup}^M"
+screen -S "$session_name" -X detachs

scripts-pretrain-resnet50/stl10.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+gpu=0
+dataset=stl10
+arch=resnet50
+feature_dim=4096
+is_mixup=true # true, false
+batch_size=256
+epochs=2000
+lr=0.01
+lr_shed=cosine # step, cosine
+mixup_loss_scale=2.0 # scale w.r.t. lambda: 0.0078125 * 5 = 0.0390625
+lmbda=0.0078125
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-pretrain-resnet50/tinyimagenet.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+gpu=0
+dataset=tiny_imagenet
+arch=resnet50
+feature_dim=4096
+is_mixup=false # true, false
+batch_size=256
+epochs=2000
+lr=0.01
+lr_shed=cosine # step, cosine
+mixup_loss_scale=4.0 # scale w.r.t. lambda
+lmbda=$(echo "scale=7; 1 / ${feature_dim}" | bc)
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python main.py --lmbda ${lmbda} --corr_zero --batch_size ${batch_size} --feature_dim ${feature_dim} --dataset ${dataset} --is_mixup ${is_mixup} --mixup_loss_scale ${mixup_loss_scale} --epochs ${epochs} --arch ${arch} --gpu ${gpu} --lr_shed ${lr_shed} --lr ${lr}^M"
+screen -S "$session_name" -X detach

scripts-transfer-resnet18/cifar10-to-x.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+gpu=0
+dataset=cifar10
+arch=resnet18
+batch_size=128
+wandb_group='best-mbt'
+model_path=checkpoints/4wdhbpcf_0.0078125_1024_256_cifar10_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+transfer_dataset='dtd'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='mnist'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='fashionmnist'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='cu_birds'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='vgg_flower'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='traffic_sign'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='aircraft'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+screen -S "$session_name" -X detach

scripts-transfer-resnet18/cifar100-to-x.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+gpu=0
+dataset=cifar100
+arch=resnet18
+batch_size=128
+wandb_group='mbt'
+model_path=checkpoints/76kk7scz_0.0078125_1024_256_cifar100_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+transfer_dataset='dtd'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='mnist'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='fashionmnist'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='cu_birds'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='vgg_flower'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='traffic_sign'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='aircraft'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+screen -S "$session_name" -X detach

scripts-transfer-resnet18/stl10-to-x-bt.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+gpu=0
+dataset=stl10
+arch=resnet18
+batch_size=128
+wandb_group='mbt'
+model_path=checkpoints/i7det4xq_0.0078125_1024_256_stl10_model.pth
+
+timestamp=$(date +"%Y%m%d%H%M%S")
+session_name="python_session_$timestamp"
+echo ${session_name}
+screen -dmS "$session_name"
+screen -S "$session_name" -X stuff "conda activate ssl-aug^M"
+transfer_dataset='dtd'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='mnist'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='fashionmnist'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='cu_birds'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='vgg_flower'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='traffic_sign'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+transfer_dataset='aircraft'
+screen -S "$session_name" -X stuff "CUDA_VISIBLE_DEVICES=${gpu} python evaluate_transfer.py --dataset ${dataset} --transfer_dataset ${transfer_dataset} --model_path ${model_path} --arch ${arch} --screen ${session_name} --wandb_group ${wandb_group}^M"
+screen -S "$session_name" -X detach

setup.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+mkdir ssl-aug
+mkdir Barlow-Twins-HSIC
+mkdir /data/wbandar1/projects/ssl-aug-artifacts/results
+git clone https://github.com/wgcban/ssl-aug.git Barlow-Twins-HSIC
+
+cd Barlow-Twins-HSIC
+conda env create -f environment.yml
+conda activate ssl-aug
+
+

ssl-sota/README.md

@@ -0,0 +1,87 @@
+# Self-Supervised Representation Learning
+
+Official repository of the paper **Whitening for Self-Supervised Representation Learning** 
+
+ICML 2021 | [arXiv:2007.06346](https://arxiv.org/abs/2007.06346)
+
+It includes 3 types of losses:
+- W-MSE [arXiv](https://arxiv.org/abs/2007.06346)
+- Contrastive [SimCLR arXiv](https://arxiv.org/abs/2002.05709)
+- BYOL [arXiv](https://arxiv.org/abs/2006.07733)
+
+And 5 datasets:
+- CIFAR-10 and CIFAR-100
+- STL-10
+- Tiny ImageNet
+- ImageNet-100
+Checkpoints are stored in `data` each 100 epochs during training.
+
+The implementation is optimized for a single GPU, although multiple are also supported. It includes fast evaluation: we pre-compute embeddings for the entire dataset and then train a classifier on top. The evaluation of the ResNet-18 encoder takes about one minute.
+
+## Installation
+
+The implementation is based on PyTorch. Logging works on [wandb.ai](https://wandb.ai/). See `docker/Dockerfile`.
+
+#### ImageNet-100
+To get this dataset, take the original ImageNet and filter out [this subset of classes](https://github.com/HobbitLong/CMC/blob/master/imagenet100.txt). We do not use augmentations during testing, and loading big images with resizing on the fly is slow, so we can preprocess classifier train and test images. We recommend [mogrify](https://imagemagick.org/script/mogrify.php) for it. First, you need to resize to 256 (just like `torchvision.transforms.Resize(256)`) and then crop to 224 (like `torchvision.transforms.CenterCrop(224)`). Finally, put the original images to `train`, and resized to `clf` and `test`.
+
+## Usage
+
+Detailed settings are good by default, to see all options:
+```
+python -m train --help
+python -m test --help
+```
+
+To reproduce the results from [table 1](https://arxiv.org/abs/2007.06346):
+#### W-MSE 4
+```
+python -m train --dataset cifar10 --epoch 1000 --lr 3e-3 --num_samples 4 --bs 256 --emb 64 --w_size 128
+python -m train --dataset cifar100 --epoch 1000 --lr 3e-3 --num_samples 4 --bs 256 --emb 64 --w_size 128
+python -m train --dataset stl10 --epoch 2000 --lr 2e-3 --num_samples 4 --bs 256 --emb 128 --w_size 256
+python -m train --dataset tiny_in --epoch 1000 --lr 2e-3 --num_samples 4 --bs 256 --emb 128 --w_size 256
+```
+
+#### W-MSE 2
+```
+python -m train --dataset cifar10 --epoch 1000 --lr 3e-3 --emb 64 --w_size 128
+python -m train --dataset cifar100 --epoch 1000 --lr 3e-3 --emb 64 --w_size 128
+python -m train --dataset stl10 --epoch 2000 --lr 2e-3 --emb 128 --w_size 256 --w_iter 4
+python -m train --dataset tiny_in --epoch 1000 --lr 2e-3 --emb 128 --w_size 256 --w_iter 4
+```
+
+#### Contrastive
+```
+python -m train --dataset cifar10 --epoch 1000 --lr 3e-3 --emb 64 --method contrastive --arch resnet50
+python -m train --dataset cifar100 --epoch 1000 --lr 3e-3 --emb 64 --method contrastive --arch resnet50
+python -m train --dataset stl10 --epoch 2000 --lr 2e-3 --emb 128 --method contrastive --arch resnet50
+python -m train --dataset tiny_in --epoch 1000 --lr 2e-3 --emb 128 --method contrastive --arch resnet50
+```
+
+#### BYOL
+```
+python -m train --dataset cifar10 --epoch 1000 --lr 3e-3 --emb 64 --method byol
+python -m train --dataset cifar100 --epoch 1000 --lr 3e-3 --emb 64 --method byol
+python -m train --dataset stl10 --epoch 2000 --lr 2e-3 --emb 128 --method byol
+python -m train --dataset tiny_in --epoch 1000 --lr 2e-3 --emb 128 --method byol
+```
+
+#### ImageNet-100
+```
+python -m train --dataset imagenet --epoch 240 --lr 2e-3 --emb 128 --w_size 256 --crop_s0 0.08 --cj0 0.8 --cj1 0.8 --cj2 0.8 --cj3 0.2 --gs_p 0.2
+python -m train --dataset imagenet --epoch 240 --lr 2e-3 --num_samples 4 --bs 256 --emb 128 --w_size 256 --crop_s0 0.08 --cj0 0.8 --cj1 0.8 --cj2 0.8 --cj3 0.2 --gs_p 0.2
+```
+
+Use `--no_norm` to disable normalization (for Euclidean distance).
+
+## Citation
+```
+@inproceedings{ermolov2021whitening,
+  title={Whitening for self-supervised representation learning},
+  author={Ermolov, Aleksandr and Siarohin, Aliaksandr and Sangineto, Enver and Sebe, Nicu},
+  booktitle={International Conference on Machine Learning},
+  pages={3015--3024},
+  year={2021},
+  organization={PMLR}
+}
+```

ssl-sota/cfg.py

@@ -0,0 +1,152 @@
+from functools import partial
+import argparse
+from torchvision import models
+import multiprocessing
+from datasets import DS_LIST
+from methods import METHOD_LIST
+
+
+def get_cfg():
+    """ generates configuration from user input in console """
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "--method", type=str, choices=METHOD_LIST, default="w_mse", help="loss type",
+    )
+    parser.add_argument(
+        "--wandb",
+        type=str,
+        default="ssl-sota",
+        help="name of the project for logging at https://wandb.ai",
+    )
+    parser.add_argument(
+        "--byol_tau", type=float, default=0.99, help="starting tau for byol loss"
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=2,
+        help="number of samples (d) generated from each image",
+    )
+
+    addf = partial(parser.add_argument, type=float)
+    addf("--cj0", default=0.4, help="color jitter brightness")
+    addf("--cj1", default=0.4, help="color jitter contrast")
+    addf("--cj2", default=0.4, help="color jitter saturation")
+    addf("--cj3", default=0.1, help="color jitter hue")
+    addf("--cj_p", default=0.8, help="color jitter probability")
+    addf("--gs_p", default=0.1, help="grayscale probability")
+    addf("--crop_s0", default=0.2, help="crop size from")
+    addf("--crop_s1", default=1.0, help="crop size to")
+    addf("--crop_r0", default=0.75, help="crop ratio from")
+    addf("--crop_r1", default=(4 / 3), help="crop ratio to")
+    addf("--hf_p", default=0.5, help="horizontal flip probability")
+
+    parser.add_argument(
+        "--no_lr_warmup",
+        dest="lr_warmup",
+        action="store_false",
+        help="do not use learning rate warmup",
+    )
+    parser.add_argument(
+        "--no_add_bn", dest="add_bn", action="store_false", help="do not use BN in head"
+    )
+    parser.add_argument("--knn", type=int, default=5, help="k in k-nn classifier")
+    parser.add_argument("--fname", type=str, help="load model from file")
+    parser.add_argument(
+        "--lr_step",
+        type=str,
+        choices=["cos", "step", "none"],
+        default="step",
+        help="learning rate schedule type",
+    )
+    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
+    parser.add_argument(
+        "--eta_min", type=float, default=0, help="min learning rate (for --lr_step cos)"
+    )
+    parser.add_argument(
+        "--adam_l2", type=float, default=1e-6, help="weight decay (L2 penalty)"
+    )
+    parser.add_argument("--T0", type=int, help="period (for --lr_step cos)")
+    parser.add_argument(
+        "--Tmult", type=int, default=1, help="period factor (for --lr_step cos)"
+    )
+    parser.add_argument(
+        "--w_eps", type=float, default=1e-4, help="eps for stability for whitening"
+    )
+    parser.add_argument(
+        "--head_layers", type=int, default=2, help="number of FC layers in head"
+    )
+    parser.add_argument(
+        "--head_size", type=int, default=1024, help="size of FC layers in head"
+    )
+
+    parser.add_argument(
+        "--w_size", type=int, default=128, help="size of sub-batch for W-MSE loss"
+    )
+    parser.add_argument(
+        "--w_iter",
+        type=int,
+        default=1,
+        help="iterations for whitening matrix estimation",
+    )
+
+    parser.add_argument(
+        "--no_norm", dest="norm", action="store_false", help="don't normalize latents",
+    )
+    parser.add_argument(
+        "--tau", type=float, default=0.5, help="contrastive loss temperature"
+    )
+
+    parser.add_argument("--epoch", type=int, default=200, help="total epoch number")
+    parser.add_argument(
+        "--eval_every_drop",
+        type=int,
+        default=5,
+        help="how often to evaluate after learning rate drop",
+    )
+    parser.add_argument(
+        "--eval_every", type=int, default=20, help="how often to evaluate"
+    )
+    parser.add_argument("--emb", type=int, default=64, help="embedding size")
+    parser.add_argument(
+        "--bs", type=int, default=384, help="number of original images in batch N",
+    )
+    parser.add_argument(
+        "--drop",
+        type=int,
+        nargs="*",
+        default=[50, 25],
+        help="milestones for learning rate decay (0 = last epoch)",
+    )
+    parser.add_argument(
+        "--drop_gamma",
+        type=float,
+        default=0.2,
+        help="multiplicative factor of learning rate decay",
+    )
+    parser.add_argument(
+        "--arch",
+        type=str,
+        choices=[x for x in dir(models) if "resn" in x],
+        default="resnet18",
+        help="encoder architecture",
+    )
+    parser.add_argument("--dataset", type=str, choices=DS_LIST, default="cifar10")
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=0,
+        help="dataset workers number",
+    )
+    parser.add_argument(
+        "--clf",
+        type=str,
+        default="sgd",
+        choices=["sgd", "knn", "lbfgs"],
+        help="classifier for test.py",
+    )
+    parser.add_argument(
+        "--eval_head", action="store_true", help="eval head output instead of model",
+    )
+    parser.add_argument("--imagenet_path", type=str, default="~/IN100/")
+    return parser.parse_args()

ssl-sota/datasets/__init__.py

@@ -0,0 +1,22 @@
+from .cifar10 import CIFAR10
+from .cifar100 import CIFAR100
+from .stl10 import STL10
+from .tiny_in import TinyImageNet
+from .imagenet import ImageNet
+
+
+DS_LIST = ["cifar10", "cifar100", "stl10", "tinyimagenet", "imagenet"]
+
+
+def get_ds(name):
+    assert name in DS_LIST
+    if name == "cifar10":
+        return CIFAR10
+    elif name == "cifar100":
+        return CIFAR100
+    elif name == "stl10":
+        return STL10
+    elif name == "tinyimagenet":
+        return TinyImageNet
+    elif name == "imagenet":
+        return ImageNet

ssl-sota/datasets/base.py

@@ -0,0 +1,67 @@
+from abc import ABCMeta, abstractmethod
+from functools import lru_cache
+from torch.utils.data import DataLoader
+
+
+class BaseDataset(metaclass=ABCMeta):
+    """
+        base class for datasets, it includes 3 types:
+            - for self-supervised training,
+            - for classifier training for evaluation,
+            - for testing
+    """
+
+    def __init__(
+        self, bs_train, aug_cfg, num_workers, bs_clf=1000, bs_test=1000,
+    ):
+        self.aug_cfg = aug_cfg
+        self.bs_train, self.bs_clf, self.bs_test = bs_train, bs_clf, bs_test
+        self.num_workers = num_workers
+
+    @abstractmethod
+    def ds_train(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def ds_clf(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def ds_test(self):
+        raise NotImplementedError
+
+    @property
+    @lru_cache()
+    def train(self):
+        return DataLoader(
+            dataset=self.ds_train(),
+            batch_size=self.bs_train,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            drop_last=True,
+        )
+
+    @property
+    @lru_cache()
+    def clf(self):
+        return DataLoader(
+            dataset=self.ds_clf(),
+            batch_size=self.bs_clf,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            drop_last=True,
+        )
+
+    @property
+    @lru_cache()
+    def test(self):
+        return DataLoader(
+            dataset=self.ds_test(),
+            batch_size=self.bs_test,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            drop_last=False,
+        )

ssl-sota/datasets/cifar10.py

@@ -0,0 +1,26 @@
+from torchvision.datasets import CIFAR10 as C10
+import torchvision.transforms as T
+from .transforms import MultiSample, aug_transform
+from .base import BaseDataset
+
+
+def base_transform():
+    return T.Compose(
+        [T.ToTensor(), T.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))]
+    )
+
+
+class CIFAR10(BaseDataset):
+    def ds_train(self):
+        t = MultiSample(
+            aug_transform(32, base_transform, self.aug_cfg), n=self.aug_cfg.num_samples
+        )
+        return C10(root="/mnt/store/wbandar1/datasets/", train=True, download=True, transform=t)
+
+    def ds_clf(self):
+        t = base_transform()
+        return C10(root="/mnt/store/wbandar1/datasets/", train=True, download=True, transform=t)
+
+    def ds_test(self):
+        t = base_transform()
+        return C10(root="/mnt/store/wbandar1/datasets/", train=False, download=True, transform=t)