Added VicReg - Variance-Invariance-Covariance Regularization for SSL

examples/vision/barlow_twins.py

@@ -2,15 +2,14 @@
 Title: Barlow Twins for Contrastive SSL
 Author: [Abhiraam Eranti](https://github.com/dewball345)
 Date created: 11/4/21
-Last modified: 12/20/21
+Last modified: 7/16/24
 Description: A keras implementation of Barlow Twins (constrastive SSL with redundancy reduction).
 Accelerator: GPU
 """
+
 """
 ## Introduction
-"""
 
-"""
 Self-supervised learning (SSL) is a relatively novel technique in which a model
 learns from unlabeled data, and is often used when the data is corrupted or
 if there is very little of it. A practical use for SSL is to create
@@ -51,7 +50,7 @@
 Also, it is simpler than other methods.
 
 This notebook can train a Barlow Twins model and reach up to
-64% validation accuracy on the CIFAR-10 dataset.
+67% validation accuracy on the CIFAR-10 dataset.
 """
 
 """
@@ -61,13 +60,11 @@
 
 
 
-
 """
 
 """
 ### High-Level Theory
 
-
 """
 
 """
@@ -127,7 +124,6 @@
 Original Implementation:
  [facebookresearch/barlowtwins](https://github.com/facebookresearch/barlowtwins)
 
-
 """
 
 """
@@ -136,6 +132,7 @@
 
 """shell
 pip install tensorflow-addons
+pip install --upgrade-strategy=only-if-needed tensorflow_similarity
 """
 
 import os
@@ -153,6 +150,7 @@
 import numpy as np  # np.random.random
 import matplotlib.pyplot as plt  # graphs
 import datetime  # tensorboard logs naming
+import tensorflow_similarity  # loss function module
 
 # XLA optimization for faster performance(up to 10-15 minutes total time saved)
 tf.config.optimizer.set_jit(True)
@@ -605,9 +603,15 @@ def plot_values(batch: tuple):
 
 After this the two parts are summed together.
 
+We will be using the 
+[BarlowLoss](https://github.com/tensorflow/similarity/blob/master/tensorflow_similarity/losses/barlow.py)
+module from Tensorflow Similarity
 
+A from-scratch implementation is also included below.
+"""
 
-
+"""
+### From-Scratch implementation (for understanding purposes)
 """
 
 
@@ -767,18 +771,17 @@ class to make the cross corr matrix, then finds the loss and
 """
 
 """
-Resnet encoder network implementation:
+### Resnet encoder network implementation:
 """
 
 
 class ResNet34:
     """Resnet34 class.
 
-        Responsible for the Resnet 34 architecture.
+    Responsible for the Resnet 34 architecture.
     Modified from
     https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2.
-    https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2.
-        View their website for more information.
+            View their website for more information.
     """
 
     def identity_block(self, x, filter):
@@ -846,7 +849,7 @@ def __call__(self, shape=(32, 32, 3)):
 
 
 """
-Projector network:
+### Projector network:
 """
 
 
@@ -960,7 +963,9 @@ def train_step(self, batch: tf.Tensor) -> tf.Tensor:
 # chose the LAMB optimizer due to high batch sizes. Converged MUCH faster
 # than ADAM or SGD
 optimizer = tfa.optimizers.LAMB()
-loss = BarlowLoss(BATCH_SIZE)
+
+# We can just drop in either one of the two(results will be similar for both)
+loss = tensorflow_similarity.losses.Barlow()  # BarlowLoss(BATCH_SIZE)
 
 bm.compile(optimizer=optimizer, loss=loss)
 
@@ -1023,12 +1028,12 @@ def train_step(self, batch: tf.Tensor) -> tf.Tensor:
 
 *   Barlow Twins is a simple and concise method for contrastive and self-supervised
 learning.
-*   With this resnet-34 model architecture, we were able to reach 62-64% validation
+*   With this resnet-34 model architecture, we were able to reach 67% validation
 accuracy.
 
 ## Use-Cases of Barlow-Twins(and contrastive learning in General)
 
-*   Semi-supervised learning: You can see that this model gave a 62-64% boost in accuracy
+*   Semi-supervised learning: You can see that this model gave a 67% boost in accuracy
 when it wasn't even trained with the labels. It can be used when you have little labeled
 data but a lot of unlabeled data.
 * You do barlow twins training on the unlabeled data, and then you do secondary training
@@ -1045,5 +1050,4 @@ def train_step(self, batch: tf.Tensor) -> tf.Tensor:
   * Thanks to Yashowardhan Shinde for writing the article.
 
 
-
 """

examples/vision/ipynb/barlow_twins.ipynb

@@ -69,7 +69,7 @@
     "Also, it is simpler than other methods.\n",
     "\n",
     "This notebook can train a Barlow Twins model and reach up to\n",
-    "64% validation accuracy on the CIFAR-10 dataset."
+    "67% validation accuracy on the CIFAR-10 dataset."
    ]
   },
   {
@@ -82,7 +82,6 @@
     "\n",
     "\n",
     "\n",
-    "\n",
     ""
    ]
   },
@@ -92,8 +91,7 @@
     "colab_type": "text"
    },
    "source": [
-    "### High-Level Theory\n",
-    ""
+    "### High-Level Theory"
    ]
   },
   {
@@ -171,8 +169,7 @@
     "Reduction](https://arxiv.org/abs/2103.03230)\n",
     "\n",
     "Original Implementation:\n",
-    " [facebookresearch/barlowtwins](https://github.com/facebookresearch/barlowtwins)\n",
-    ""
+    " [facebookresearch/barlowtwins](https://github.com/facebookresearch/barlowtwins)"
    ]
   },
   {
@@ -192,7 +189,8 @@
    },
    "outputs": [],
    "source": [
-    "!pip install tensorflow-addons"
+    "!!pip install tensorflow-addons\n",
+    "!!pip install --upgrade-strategy=only-if-needed tensorflow_similarity"
    ]
   },
   {
@@ -218,6 +216,7 @@
     "import numpy as np  # np.random.random\n",
     "import matplotlib.pyplot as plt  # graphs\n",
     "import datetime  # tensorboard logs naming\n",
+    "import tensorflow_similarity  # loss function module\n",
     "\n",
     "# XLA optimization for faster performance(up to 10-15 minutes total time saved)\n",
     "tf.config.optimizer.set_jit(True)"
@@ -761,8 +760,20 @@
     "\n",
     "After this the two parts are summed together.\n",
     "\n",
+    "We will be using the [BarlowLoss](https://github.com/tensorflow/similarity/blob/master/tensorflow_similarity/lo\n",
+    "sses/barlow.py)\n",
+    "module from Tensorflow Similarity\n",
     "\n",
-    ""
+    "A from-scratch implementation is also included below."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "### From-Scratch implementation(for understanding purposes)"
    ]
   },
   {
@@ -941,7 +952,7 @@
     "colab_type": "text"
    },
    "source": [
-    "Resnet encoder network implementation:"
+    "### Resnet encoder network implementation:"
    ]
   },
   {
@@ -956,11 +967,13 @@
     "class ResNet34:\n",
     "    \"\"\"Resnet34 class.\n",
     "\n",
-    "        Responsible for the Resnet 34 architecture.\n",
-    "    Modified from\n",
+    "            Responsible for the Resnet 34 architecture.\n",
+    "        Modified from\n",
+    "    https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2.\n",
+    "    https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2.\n",
     "    https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2.\n",
     "    https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2.\n",
-    "        View their website for more information.\n",
+    "            View their website for more information.\n",
     "    \"\"\"\n",
     "\n",
     "    def identity_block(self, x, filter):\n",
@@ -1034,7 +1047,7 @@
     "colab_type": "text"
    },
    "source": [
-    "Projector network:"
+    "### Projector network:"
    ]
   },
   {
@@ -1184,7 +1197,9 @@
     "# chose the LAMB optimizer due to high batch sizes. Converged MUCH faster\n",
     "# than ADAM or SGD\n",
     "optimizer = tfa.optimizers.LAMB()\n",
-    "loss = BarlowLoss(BATCH_SIZE)\n",
+    "\n",
+    "# We can just drop in either one of the two(results will be similar for both)\n",
+    "loss = tensorflow_similarity.losses.Barlow()  # BarlowLoss(BATCH_SIZE)\n",
     "\n",
     "bm.compile(optimizer=optimizer, loss=loss)\n",
     "\n",
@@ -1267,12 +1282,12 @@
     "\n",
     "*   Barlow Twins is a simple and concise method for contrastive and self-supervised\n",
     "learning.\n",
-    "*   With this resnet-34 model architecture, we were able to reach 62-64% validation\n",
+    "*   With this resnet-34 model architecture, we were able to reach 67% validation\n",
     "accuracy.\n",
     "\n",
     "## Use-Cases of Barlow-Twins(and contrastive learning in General)\n",
     "\n",
-    "*   Semi-supervised learning: You can see that this model gave a 62-64% boost in accuracy\n",
+    "*   Semi-supervised learning: You can see that this model gave a 67% boost in accuracy\n",
     "when it wasn't even trained with the labels. It can be used when you have little labeled\n",
     "data but a lot of unlabeled data.\n",
     "* You do barlow twins training on the unlabeled data, and then you do secondary training\n",
@@ -1284,12 +1299,13 @@
     "* [Original Pytorch Implementation](https://github.com/facebookresearch/barlowtwins)\n",
     "* [Sayak Paul's\n",
     "Implementation](https://colab.research.google.com/github/sayakpaul/Barlow-Twins-TF/blob/main/Barlow_Twins.ipynb#scrollTo=GlWepkM8_prl).\n",
+    "Implementation](https://colab.research.google.com/github/sayakpaul/Barlow-Twins-TF/blob/main/Barlow_Twins.ipynb#scrollTo=GlWepkM8_prl).\n",
     "* Thanks to Sayak Paul for his implementation. It helped me with debugging and\n",
     "comparisons of accuracy, loss.\n",
     "* [resnet34\n",
     "implementation](https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2)\n",
+    "implementation](https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2)\n",
     "  * Thanks to Yashowardhan Shinde for writing the article.\n",
-    "\n",
     ""
    ]
   }