SimCLR（Simple Framework for Contrastive Learning of Representations）是一種學(xué)習(xí)圖像表示的自監(jiān)督技術(shù)。 與傳統(tǒng)的監(jiān)督學(xué)習(xí)方法不同，SimCLR 不依賴(lài)標(biāo)記數(shù)據(jù)來(lái)學(xué)習(xí)有用的表示。 它利用對(duì)比學(xué)習(xí)框架來(lái)學(xué)習(xí)一組有用的特征，這些特征可以從未標(biāo)記的圖像中捕獲高級(jí)語(yǔ)義信息。

(相關(guān)資料圖)

SimCLR 已被證明在各種圖像分類(lèi)基準(zhǔn)上優(yōu)于最先進(jìn)的無(wú)監(jiān)督學(xué)習(xí)方法。 并且它學(xué)習(xí)到的表示可以很容易地轉(zhuǎn)移到下游任務(wù)，例如對(duì)象檢測(cè)、語(yǔ)義分割和小樣本學(xué)習(xí)，只需在較小的標(biāo)記數(shù)據(jù)集上進(jìn)行最少的微調(diào)。

SimCLR 主要思想是通過(guò)增強(qiáng)模塊 T 將圖像與同一圖像的其他增強(qiáng)版本進(jìn)行對(duì)比，從而學(xué)習(xí)圖像的良好表示。這是通過(guò)通過(guò)編碼器網(wǎng)絡(luò) f(.) 映射圖像，然后進(jìn)行投影來(lái)完成的。 head g(.) 將學(xué)習(xí)到的特征映射到低維空間。 然后在同一圖像的兩個(gè)增強(qiáng)版本的表示之間計(jì)算對(duì)比損失，以鼓勵(lì)對(duì)同一圖像的相似表示和對(duì)不同圖像的不同表示。

本文我們將深入研究 SimCLR 框架并探索該算法的關(guān)鍵組件，包括數(shù)據(jù)增強(qiáng)、對(duì)比損失函數(shù)以及編碼器和投影的head 架構(gòu)。

我們這里使用來(lái)自 Kaggle 的垃圾分類(lèi)數(shù)據(jù)集來(lái)進(jìn)行實(shí)驗(yàn)

增強(qiáng)模塊

SimCLR 中最重要的就是轉(zhuǎn)換圖像的增強(qiáng)模塊。 SimCLR 論文的作者建議，強(qiáng)大的數(shù)據(jù)增強(qiáng)對(duì)于無(wú)監(jiān)督學(xué)習(xí)很有用。 因此，我們將遵循論文中推薦的方法。

調(diào)整大小的隨機(jī)裁剪50% 概率的隨機(jī)水平翻轉(zhuǎn)隨機(jī)顏色失真（顏色抖動(dòng)概率為 80%，顏色下降概率為 20%）50% 概率為隨機(jī)高斯模糊

def get_complete_transform(output_shape, kernel_size, s=1.0):     """    Color distortion transform         Args:        s: Strength parameter             Returns:        A color distortion transform    """     rnd_crop = RandomResizedCrop(output_shape)     rnd_flip = RandomHorizontalFlip(p=0.5)          color_jitter = ColorJitter(0.8*s, 0.8*s, 0.8*s, 0.2*s)     rnd_color_jitter = RandomApply([color_jitter], p=0.8)          rnd_gray = RandomGrayscale(p=0.2)     gaussian_blur = GaussianBlur(kernel_size=kernel_size)     rnd_gaussian_blur = RandomApply([gaussian_blur], p=0.5)     to_tensor = ToTensor()     image_transform = Compose([         to_tensor,         rnd_crop,         rnd_flip,         rnd_color_jitter,         rnd_gray,         rnd_gaussian_blur,    ])     return image_transform  class ContrastiveLearningViewGenerator(object):     """    Take 2 random crops of 1 image as the query and key.    """     def __init__(self, base_transform, n_views=2):         self.base_transform = base_transform         self.n_views = n_views              def __call__(self, x):         views = [self.base_transform(x) for i in range(self.n_views)]         return views

下一步就是定義一個(gè)PyTorch 的 Dataset 。

class CustomDataset(Dataset):     def __init__(self, list_images, transform=None):         """        Args:            list_images (list): List of all the images            transform (callable, optional): Optional transform to be applied on a sample.        """         self.list_images = list_images         self.transform = transform              def __len__(self):         return len(self.list_images)          def __getitem__(self, idx):         if torch.is_tensor(idx):             idx = idx.tolist()                      img_name = self.list_images[idx]         image = io.imread(img_name)         if self.transform:             image = self.transform(image)                      return image

作為樣例，我們使用比較小的模型 ResNet18 作為主干，所以他的輸入是 224x224 圖像，我們按照要求設(shè)置一些參數(shù)并生成dataloader

out_shape = [224, 224] kernel_size = [21, 21] # 10% of out_shape  # Custom transform base_transforms = get_complete_transform(output_shape=out_shape, kernel_size=kernel_size, s=1.0) custom_transform = ContrastiveLearningViewGenerator(base_transform=base_transforms)  garbage_ds = CustomDataset(     list_images=glob.glob("/kaggle/input/garbage-classification/garbage_classification/*/*.jpg"),     transform=custom_transform )  BATCH_SZ = 128  # Build DataLoader train_dl = torch.utils.data.DataLoader(     garbage_ds,     batch_size=BATCH_SZ,     shuffle=True,     drop_last=True,     pin_memory=True)

SimCLR

我們已經(jīng)準(zhǔn)備好了數(shù)據(jù)，開(kāi)始對(duì)模型進(jìn)行復(fù)現(xiàn)。上面的增強(qiáng)模塊提供了圖像的兩個(gè)增強(qiáng)視圖，它們通過(guò)編碼器前向傳遞以獲得相應(yīng)的表示。 SimCLR 的目標(biāo)是通過(guò)鼓勵(lì)模型從兩個(gè)不同的增強(qiáng)視圖中學(xué)習(xí)對(duì)象的一般表示來(lái)最大化這些不同學(xué)習(xí)表示之間的相似性。

編碼器網(wǎng)絡(luò)的選擇不受限制，可以是任何架構(gòu)。 上面已經(jīng)說(shuō)了，為了簡(jiǎn)單演示，我們使用 ResNet18。 編碼器模型學(xué)習(xí)到的表示決定了相似性系數(shù)，為了提高這些表示的質(zhì)量，SimCLR 使用投影頭將編碼向量投影到更豐富的潛在空間中。 這里我們將ResNet18的512維度的特征投影到256的空間中，看著很復(fù)雜，其實(shí)就是加了一個(gè)帶relu的mlp。

class Identity(nn.Module):     def __init__(self):         super(Identity, self).__init__()     def forward(self, x):         return x      class SimCLR(nn.Module):     def __init__(self, linear_eval=False):         super().__init__()         self.linear_eval = linear_eval         resnet18 = models.resnet18(pretrained=False)         resnet18.fc = Identity()         self.encoder = resnet18         self.projection = nn.Sequential(             nn.Linear(512, 512),             nn.ReLU(),             nn.Linear(512, 256)        )     def forward(self, x):         if not self.linear_eval:             x = torch.cat(x, dim=0)         encoding = self.encoder(x)         projection = self.projection(encoding)         return projection

對(duì)比損失

對(duì)比損失函數(shù)，也稱(chēng)為歸一化溫度標(biāo)度交叉熵?fù)p失 (NT-Xent)，是 SimCLR 的一個(gè)關(guān)鍵組成部分，它鼓勵(lì)模型學(xué)習(xí)相同圖像的相似表示和不同圖像的不同表示。

NT-Xent 損失是使用一對(duì)通過(guò)編碼器網(wǎng)絡(luò)傳遞的圖像的增強(qiáng)視圖來(lái)計(jì)算的，以獲得它們相應(yīng)的表示。 對(duì)比損失的目標(biāo)是鼓勵(lì)同一圖像的兩個(gè)增強(qiáng)視圖的表示相似，同時(shí)迫使不同圖像的表示不相似。

NT-Xent 將 softmax 函數(shù)應(yīng)用于增強(qiáng)視圖表示的成對(duì)相似性。 softmax 函數(shù)應(yīng)用于小批量?jī)?nèi)的所有表示對(duì)，得到每個(gè)圖像的相似性概率分布。 溫度參數(shù)temperature 用于在應(yīng)用 softmax 函數(shù)之前縮放成對(duì)相似性，這有助于在優(yōu)化過(guò)程中獲得更好的梯度。

在獲得相似性的概率分布后，通過(guò)最大化同一圖像的匹配表示的對(duì)數(shù)似然和最小化不同圖像的不匹配表示的對(duì)數(shù)似然來(lái)計(jì)算 NT-Xent 損失。

LABELS = torch.cat([torch.arange(BATCH_SZ) for i in range(2)], dim=0) LABELS = (LABELS.unsqueeze(0) == LABELS.unsqueeze(1)).float() #one-hot representations LABELS = LABELS.to(DEVICE)  def ntxent_loss(features, temp):     """    NT-Xent Loss.         Args:        z1: The learned representations from first branch of projection head        z2: The learned representations from second branch of projection head    Returns:        Loss    """     similarity_matrix = torch.matmul(features, features.T)     mask = torch.eye(LABELS.shape[0], dtype=torch.bool).to(DEVICE)     labels = LABELS[~mask].view(LABELS.shape[0], -1)     similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)          positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)          negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)          logits = torch.cat([positives, negatives], dim=1)     labels = torch.zeros(logits.shape[0], dtype=torch.long).to(DEVICE)          logits = logits / temp     return logits, labels

所有的準(zhǔn)備都完成了，讓我們訓(xùn)練 SimCLR 看看效果！

simclr_model = SimCLR().to(DEVICE) criterion = nn.CrossEntropyLoss().to(DEVICE) optimizer = torch.optim.Adam(simclr_model.parameters())  epochs = 10 with tqdm(total=epochs) as pbar:     for epoch in range(epochs):         t0 = time.time()         running_loss = 0.0         for i, views in enumerate(train_dl):             projections = simclr_model([view.to(DEVICE) for view in views])             logits, labels = ntxent_loss(projections, temp=2)             loss = criterion(logits, labels)             optimizer.zero_grad()             loss.backward()             optimizer.step()                          # print stats             running_loss += loss.item()             if i%10 == 9: # print every 10 mini-batches                 print(f"Epoch: {epoch+1} Batch: {i+1} Loss: {(running_loss/100):.4f}")                 running_loss = 0.0         pbar.update(1)         print(f"Time taken: {((time.time()-t0)/60):.3f} mins")

上面代碼訓(xùn)練了10輪，假設(shè)我們已經(jīng)完成了預(yù)訓(xùn)練過(guò)程，可以將預(yù)訓(xùn)練的編碼器用于我們想要的下游任務(wù)。這可以通過(guò)下面的代碼來(lái)完成。

from torchvision.transforms import Resize, CenterCrop resize = Resize(255) ccrop = CenterCrop(224) ttensor = ToTensor()  custom_transform = Compose([     resize,     ccrop,     ttensor, ])  garbage_ds = ImageFolder(     root="/kaggle/input/garbage-classification/garbage_classification/",     transform=custom_transform )  classes = len(garbage_ds.classes)  BATCH_SZ = 128  train_dl = torch.utils.data.DataLoader(     garbage_ds,     batch_size=BATCH_SZ,     shuffle=True,     drop_last=True,     pin_memory=True, )  class Identity(nn.Module):     def __init__(self):         super(Identity, self).__init__()     def forward(self, x):         return x      class LinearEvaluation(nn.Module):     def __init__(self, model, classes):         super().__init__()         simclr = model         simclr.linear_eval=True         simclr.projection = Identity()         self.simclr = simclr         for param in self.simclr.parameters():             param.requires_grad = False         self.linear = nn.Linear(512, classes)     def forward(self, x):         encoding = self.simclr(x)         pred = self.linear(encoding)         return pred        eval_model = LinearEvaluation(simclr_model, classes).to(DEVICE) criterion = nn.CrossEntropyLoss().to(DEVICE) optimizer = torch.optim.Adam(eval_model.parameters())  preds, labels = [], [] correct, total = 0, 0  with torch.no_grad():     t0 = time.time()     for img, gt in tqdm(train_dl):         image = img.to(DEVICE)         label = gt.to(DEVICE)         pred = eval_model(image)         _, pred = torch.max(pred.data, 1)         total += label.size(0)         correct += (pred == label).float().sum().item()      print(f"Time taken: {((time.time()-t0)/60):.3f} mins")      print(         "Accuracy of the network on the {} Train images: {} %".format(             total, 100 * correct / total        )    )

上面的代碼最主要的部分就是讀取剛剛訓(xùn)練的simclr模型，然后凍結(jié)所有的權(quán)重，然后再創(chuàng)建一個(gè)分類(lèi)頭self.linear ，進(jìn)行下游的分類(lèi)任務(wù)

總結(jié)

本文介紹了SimCLR框架，并使用它來(lái)預(yù)訓(xùn)練隨機(jī)初始化權(quán)重的ResNet18。預(yù)訓(xùn)練是深度學(xué)習(xí)中使用的一種強(qiáng)大的技術(shù)，用于在大型數(shù)據(jù)集上訓(xùn)練模型，學(xué)習(xí)可以轉(zhuǎn)移到其他任務(wù)中的有用特征。SimCLR論文認(rèn)為，批量越大，性能越好。我們的實(shí)現(xiàn)只使用128個(gè)批大小，只訓(xùn)練10個(gè)epoch。所以這不是模型的最佳性能，如果需要性能對(duì)比還需要進(jìn)一步的訓(xùn)練。

下圖是論文作者給出的性能結(jié)論：

關(guān)鍵詞：