![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/13-contrastive-learning/simclr_data_augmentations.jpg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/13-contrastive-learning/crop_views.svg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/13-contrastive-learning/simclr_network_setup.svg\")
\n", "After we have trained our model via contrastive learning, we can deploy it on downstream tasks and see how well it performs with little data.\n", "A common setup, which also verifies whether the model has learned generalized representations, is to perform Logistic Regression on the features.\n", "In other words, we learn a single, linear layer that maps the representations to a class prediction.\n", "Since the base network $f(\\cdot)$ is not changed during the training process, the model can only perform well if the representations of $h$ describe all features that might be necessary for the task.\n", "Further, we do not have to worry too much about overfitting since we have very few parameters that are trained.\n", "Hence, we might expect that the model can perform well even with very little data.\n", "\n", "First, let's implement a simple Logistic Regression setup for which we assume that the images already have been encoded in their feature vectors.\n", "If very little data is available, it might be beneficial to dynamically encode the images during training so that we can also apply data augmentations.\n", "However, the way we implement it here is much more efficient and can be trained within a few seconds.\n", "Further, using data augmentations did not show any significant gain in this simple setup."]}, {"cell_type": "code", "execution_count": 12, "id": "5007f0d1", "metadata": {"execution": {"iopub.execute_input": "2025-05-01T10:51:49.118022Z", "iopub.status.busy": "2025-05-01T10:51:49.116935Z", "iopub.status.idle": "2025-05-01T10:51:49.129039Z", "shell.execute_reply": "2025-05-01T10:51:49.127993Z"}, "papermill": {"duration": 0.073307, "end_time": "2025-05-01T10:51:49.130428", "exception": false, "start_time": "2025-05-01T10:51:49.057121", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class LogisticRegression(pl.LightningModule):\n", " def __init__(self, feature_dim, num_classes, lr, weight_decay, max_epochs=100):\n", " super().__init__()\n", " self.save_hyperparameters()\n", " # Mapping from representation h to classes\n", " self.model = nn.Linear(feature_dim, num_classes)\n", "\n", " def configure_optimizers(self):\n", " optimizer = optim.AdamW(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay)\n", " lr_scheduler = optim.lr_scheduler.MultiStepLR(\n", " optimizer, milestones=[int(self.hparams.max_epochs * 0.6), int(self.hparams.max_epochs * 0.8)], gamma=0.1\n", " )\n", " return [optimizer], [lr_scheduler]\n", "\n", " def _calculate_loss(self, batch, mode=\"train\"):\n", " feats, labels = batch\n", " preds = self.model(feats)\n", " loss = F.cross_entropy(preds, labels)\n", " acc = (preds.argmax(dim=-1) == labels).float().mean()\n", "\n", " self.log(mode + \"_loss\", loss)\n", " self.log(mode + \"_acc\", acc)\n", " return loss\n", "\n", " def training_step(self, batch, batch_idx):\n", " return self._calculate_loss(batch, mode=\"train\")\n", "\n", " def validation_step(self, batch, batch_idx):\n", " self._calculate_loss(batch, mode=\"val\")\n", "\n", " def test_step(self, batch, batch_idx):\n", " self._calculate_loss(batch, mode=\"test\")"]}, {"cell_type": "markdown", "id": "538bd05c", "metadata": {"papermill": {"duration": 0.058464, "end_time": "2025-05-01T10:51:49.247247", "exception": false, "start_time": "2025-05-01T10:51:49.188783", "status": "completed"}, "tags": []}, "source": ["The data we use is the training and test set of STL10.\n", "The training contains 500 images per class, while the test set has 800 images per class."]}, {"cell_type": "code", "execution_count": 13, "id": "ee1a6e10", "metadata": {"execution": {"iopub.execute_input": "2025-05-01T10:51:49.365246Z", "iopub.status.busy": "2025-05-01T10:51:49.364895Z", "iopub.status.idle": "2025-05-01T10:51:59.830227Z", "shell.execute_reply": "2025-05-01T10:51:59.829129Z"}, "papermill": {"duration": 10.527034, "end_time": "2025-05-01T10:51:59.832526", "exception": false, "start_time": "2025-05-01T10:51:49.305492", "status": "completed"}, "tags": []}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["Files already downloaded and verified\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Files already downloaded and verified\n", "Number of training examples: 5000\n", "Number of test examples: 8000\n"]}], "source": ["img_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])\n", "\n", "train_img_data = STL10(root=DATASET_PATH, split=\"train\", download=True, transform=img_transforms)\n", "test_img_data = STL10(root=DATASET_PATH, split=\"test\", download=True, transform=img_transforms)\n", "\n", "print(\"Number of training examples:\", len(train_img_data))\n", "print(\"Number of test examples:\", len(test_img_data))"]}, {"cell_type": "markdown", "id": "3d94238d", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.057896, "end_time": "2025-05-01T10:51:59.985646", "exception": false, "start_time": "2025-05-01T10:51:59.927750", "status": "completed"}, "tags": []}, "source": ["Next, we implement a small function to encode all images in our datasets.\n", "The output representations are then used as inputs to the Logistic Regression model."]}, {"cell_type": "code", "execution_count": 14, "id": "329e46d5", "metadata": {"execution": {"iopub.execute_input": "2025-05-01T10:52:00.103826Z", "iopub.status.busy": "2025-05-01T10:52:00.103577Z", "iopub.status.idle": "2025-05-01T10:52:00.109488Z", "shell.execute_reply": "2025-05-01T10:52:00.108669Z"}, "papermill": {"duration": 0.066661, "end_time": "2025-05-01T10:52:00.110662", "exception": false, "start_time": "2025-05-01T10:52:00.044001", "status": "completed"}, "tags": []}, "outputs": [], "source": ["@torch.no_grad()\n", "def prepare_data_features(model, dataset):\n", " # Prepare model\n", " network = deepcopy(model.convnet)\n", " network.fc = nn.Identity() # Removing projection head g(.)\n", " network.eval()\n", " network.to(device)\n", "\n", " # Encode all images\n", " data_loader = data.DataLoader(dataset, batch_size=64, num_workers=NUM_WORKERS, shuffle=False, drop_last=False)\n", " feats, labels = [], []\n", " for batch_imgs, batch_labels in tqdm(data_loader):\n", " batch_imgs = batch_imgs.to(device)\n", " batch_feats = network(batch_imgs)\n", " feats.append(batch_feats.detach().cpu())\n", " labels.append(batch_labels)\n", "\n", " feats = torch.cat(feats, dim=0)\n", " labels = torch.cat(labels, dim=0)\n", "\n", " # Sort images by labels\n", " labels, idxs = labels.sort()\n", " feats = feats[idxs]\n", "\n", " return data.TensorDataset(feats, labels)"]}, {"cell_type": "markdown", "id": "4bd938f9", "metadata": {"papermill": {"duration": 0.058285, "end_time": "2025-05-01T10:52:00.227398", "exception": false, "start_time": "2025-05-01T10:52:00.169113", "status": "completed"}, "tags": []}, "source": ["Let's apply the function to both training and test set below."]}, {"cell_type": "code", "execution_count": 15, "id": "6775136f", "metadata": {"execution": {"iopub.execute_input": "2025-05-01T10:52:00.374734Z", "iopub.status.busy": "2025-05-01T10:52:00.374377Z", "iopub.status.idle": "2025-05-01T10:52:07.878406Z", "shell.execute_reply": "2025-05-01T10:52:07.877120Z"}, "papermill": {"duration": 7.59423, "end_time": "2025-05-01T10:52:07.879798", "exception": false, "start_time": "2025-05-01T10:52:00.285568", "status": "completed"}, "tags": []}, "outputs": [{"data": {"application/vnd.jupyter.widget-view+json": {"model_id": "851e06a823de49609a7d279358814f02", "version_major": 2, "version_minor": 0}, "text/plain": [" 0%| | 0/79 [00:00\n", "\n", "\n"], "text/plain": ["