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