SimCLR（Simple Framework for Contrastive Learning of Representations）是一種學(xué)習(xí)圖像表示的自監(jiān)督技術(shù)。 與傳統(tǒng)的監(jiān)督學(xué)習(xí)方法不同，SimCLR 不依賴標(biāo)記數(shù)據(jù)來學(xué)習(xí)有用的表示。 它利用對比學(xué)習(xí)框架來學(xué)習(xí)一組有用的特征，這些特征可以從未標(biāo)記的圖像中捕獲高級語義信息。
SimCLR 已被證明在各種圖像分類基準(zhǔn)上優(yōu)于最先進(jìn)的無監(jiān)督學(xué)習(xí)方法。 并且它學(xué)習(xí)到的表示可以很容易地轉(zhuǎn)移到下游任務(wù)，例如對象檢測、語義分割和小樣本學(xué)習(xí)，只需在較小的標(biāo)記數(shù)據(jù)集上進(jìn)行最少的微調(diào)。
SimCLR 主要思想是通過增強模塊 T 將圖像與同一圖像的其他增強版本進(jìn)行對比，從而學(xué)習(xí)圖像的良好表示。這是通過通過編碼器網(wǎng)絡(luò) f(.) 映射圖像，然后進(jìn)行投影來完成的。 head g(.) 將學(xué)習(xí)到的特征映射到低維空間。 然后在同一圖像的兩個增強版本的表示之間計算對比損失，以鼓勵對同一圖像的相似表示和對不同圖像的不同表示。
本文我們將深入研究 SimCLR 框架并探索該算法的關(guān)鍵組件，包括數(shù)據(jù)增強、對比損失函數(shù)以及編碼器和投影的head 架構(gòu)。
我們這里使用來自 Kaggle 的垃圾分類數(shù)據(jù)集來進(jìn)行實驗
增強模塊
SimCLR 中最重要的就是轉(zhuǎn)換圖像的增強模塊。 SimCLR 論文的作者建議，強大的數(shù)據(jù)增強對于無監(jiān)督學(xué)習(xí)很有用。 因此，我們將遵循論文中推薦的方法。
- 調(diào)整大小的隨機裁剪
- 50% 概率的隨機水平翻轉(zhuǎn)
- 隨機顏色失真（顏色抖動概率為 80%，顏色下降概率為 20%）
- 50% 概率為隨機高斯模糊
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
下一步就是定義一個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ù)，開始對模型進(jìn)行復(fù)現(xiàn)。上面的增強模塊提供了圖像的兩個增強視圖，它們通過編碼器前向傳遞以獲得相應(yīng)的表示。 SimCLR 的目標(biāo)是通過鼓勵模型從兩個不同的增強視圖中學(xué)習(xí)對象的一般表示來最大化這些不同學(xué)習(xí)表示之間的相似性。
編碼器網(wǎng)絡(luò)的選擇不受限制，可以是任何架構(gòu)。 上面已經(jīng)說了，為了簡單演示，我們使用 ResNet18。 編碼器模型學(xué)習(xí)到的表示決定了相似性系數(shù)，為了提高這些表示的質(zhì)量，SimCLR 使用投影頭將編碼向量投影到更豐富的潛在空間中。 這里我們將ResNet18的512維度的特征投影到256的空間中，看著很復(fù)雜，其實就是加了一個帶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
對比損失
對比損失函數(shù)，也稱為歸一化溫度標(biāo)度交叉熵?fù)p失 (NT-Xent)，是 SimCLR 的一個關(guān)鍵組成部分，它鼓勵模型學(xué)習(xí)相同圖像的相似表示和不同圖像的不同表示。
NT-Xent 損失是使用一對通過編碼器網(wǎng)絡(luò)傳遞的圖像的增強視圖來計算的，以獲得它們相應(yīng)的表示。 對比損失的目標(biāo)是鼓勵同一圖像的兩個增強視圖的表示相似，同時迫使不同圖像的表示不相似。
NT-Xent 將 softmax 函數(shù)應(yīng)用于增強視圖表示的成對相似性。 softmax 函數(shù)應(yīng)用于小批量內(nèi)的所有表示對，得到每個圖像的相似性概率分布。 溫度參數(shù)temperature 用于在應(yīng)用 softmax 函數(shù)之前縮放成對相似性，這有助于在優(yōu)化過程中獲得更好的梯度。
在獲得相似性的概率分布后，通過最大化同一圖像的匹配表示的對數(shù)似然和最小化不同圖像的不匹配表示的對數(shù)似然來計算 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)練過程，可以將預(yù)訓(xùn)練的編碼器用于我們想要的下游任務(wù)。這可以通過下面的代碼來完成。
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)建一個分類頭self.linear ，進(jìn)行下游的分類任務(wù)
總結(jié)
本文介紹了SimCLR框架，并使用它來預(yù)訓(xùn)練隨機初始化權(quán)重的ResNet18。預(yù)訓(xùn)練是深度學(xué)習(xí)中使用的一種強大的技術(shù)，用于在大型數(shù)據(jù)集上訓(xùn)練模型，學(xué)習(xí)可以轉(zhuǎn)移到其他任務(wù)中的有用特征。SimCLR論文認(rèn)為，批量越大，性能越好。我們的實現(xiàn)只使用128個批大小，只訓(xùn)練10個epoch。所以這不是模型的最佳性能，如果需要性能對比還需要進(jìn)一步的訓(xùn)練。
下圖是論文作者給出的性能結(jié)論：
