Search
Duplicate

AI/ SimCLR

SimCLR(Simple Contrastive Learning of visual Representation)

SimCLR은 self-supervised 학습 방법의 하나로 constastive learning을 통해 label이 없는 이미지 데이터에서 유용한 embedding을 학습하는 방법이다.
이에 대한 아이디어는 다음과 같다.
1.
데이터셋에서 이미지를 뽑아 2가지 view로 증강 시킨다. 동일한 원본에서 변형된 것이므로 x1,x2\bold{x}_1, \bold{x}_2는 ‘의미적으로 동등한’ 쌍이다.
x∈RD→x1=t1(x),x2=t2(x)\bold{x} \in \mathbb{R}^D \to \bold{x}_1 = t_1(\bold{x}), \bold{x}_2 = t_2(\bold{x})
2.
minibatch 내의 다른 이미지를 negative 예제로 뽑는다.
x1−,...,xn−∈N(x)\bold{x}_1^-,...,\bold{x}_n^- \in N(\bold{x})
3.
어떤 feature 매핑 F:RD→REF : \mathbb{R}^D \to \mathbb{R}^E를 이용하여 (여기서 DD는 입력의 크기이고 EE는 임베딩의 크기) 각 입력 x\bold{x}에 대해 negative 와의 유사도를 최소화하면서 positive와의 유사도를 최대화한다.
J=F(t1(x))⊤F(t2(x))−log⁡∑xi−∈N(x)exp⁡[F(xi−)⊤F(t1(x))]J = F(t_1(\bold{x}))^\top F(t_2(\bold{x})) - \log \sum_{\bold{x}_i^- \in N(\bold{x})} \exp [F(\bold{x}_i^-)^\top F(t_1(\bold{x}))]
실제에서 cosine 유사도를 사용하므로 내적을 취하기 전에 FF에 의해 생성된 표현을 ℓ2\ell_2-normalization 하지만 위의 방정식에서는 생략.
SimCLR의 loss 함수는 기본적으로 다음과 같이 정의되는 InfoNCE의 loss를 따르지만, 분모 부분의 합에 positive 유사도를 제외한 형태가 된다.
L(θ;x,x+,{xk−}k=1N−1)=−log⁡exp⁡(e^θ(x)⊤e^θ(x+)/τ)exp⁡(e^θ(x)⊤e^θ(x+)/τ)+∑k=1N−1exp⁡(e^θ(x)⊤e^θ(xk−)/τ)\begin{aligned} &\mathcal{L}(\boldsymbol{\theta};\bold{x}, \bold{x}^+, \{\bold{x}_k^-\}_{k=1}^{N-1}) \\&= - \log{\exp(\hat{\bold{e}}_{\boldsymbol{\theta}}(\bold{x})^\top \hat{\bold{e}}_{\boldsymbol{\theta}}(\bold{x}^+)/\tau) \over \exp(\hat{\bold{e}}_{\boldsymbol{\theta}}(\bold{x})^\top\hat{\bold{e}}_{\boldsymbol{\theta}}(\bold{x}^+)/\tau) + \sum_{k=1}^{N-1} \exp(\hat{\bold{e}}_{\boldsymbol{\theta}}(\bold{x})^\top \hat{\bold{e}}_{\boldsymbol{\theta}}(\bold{x}_k^-)/\tau)} \end{aligned}
SimCLR은 흥미롭게도 다음의 조건부 Energy-Based Model의 형식으로 볼 수 있다.
p(x2∣x1)=exp⁡[−E(x2∣x1)]Z(x1)p(\bold{x}_2|\bold{x}_1) = {\exp[-\mathcal{E}(\bold{x}_2|\bold{x}_1)] \over Z(\bold{x}_1)}
여기서 E(x2∣x1)=−F(x2)⊤F(x1)\mathcal{E}(\bold{x}_2|\bold{x}_1) = -F(\bold{x}_2)^\top F(\bold{x}_1)은 에너지이고
Z(x)=∫exp⁡[−E(x−∣x)]dx−=∫exp⁡[F(x−)⊤F(x)]dx−Z(\bold{x}) = \int \exp [-\mathcal{E}(\bold{x}^-|\bold{x})]d\bold{x}^- = \int \exp[F(\bold{x}^-)^\top F(\bold{x})]d\bold{x}^-
는 normalization 상수(partition function이라고도 부름)이다. 이 모델의 조건부 log likelihood는 다음 형식을 갖는다.
log⁡p(x2∣x1)=F(x2)⊤F(x1)−log⁡∫exp⁡[F(x−)⊤F(x1)]dx−\log p(\bold{x}_2|\bold{x}_1) = F(\bold{x}_2)^\top F(\bold{x}_1) - \log \int \exp [F(\bold{x}^-)^\top F(\bold{x}_1)]d\bold{x}^-
앞선 방정식과 다른 점은 적분을 negative 샘플에서 유도된 Monte Carlo 상한으로 교체한다는 것이다. 따라서 contrastive learning을 조건부 에너지 기반 생성 모델의 MLE 근사로 생각할 수 있다.
SimCLR 성공의 핵심 요소(ingredient)는 데이터 증강 방법의 선택에 있다. 랜덤 cropping을 사용하여 모델이 같은 이미지의 인접 보기 예측 뿐만 아니라 전역 관점에서 지역 관점으로 강제할 수 있다. cropping 이후에 모든 이미지는 같은 크기로 resize 된다. 또한 어떤 시간 비율로 랜덤하게 flip 된다.

Sample Code

Model

간단히 ResNet18을 사용하여 SimCLR의 embedding 모델을 정의한다.
cosine 유사도 계산을 위해 모델의 forward의 마지막 단계에서 normalize를 수행하여 결과를 출력한다.
# 임베딩 모델 정의 (ResNet18 사용) class SimCLRModel(nn.Module): def __init__(self, base_model, in_features, embedding_dim): super(SimCLRModel, self).__init__() self.base_model = base_model self.base_model.fc = nn.Identity() # Fully connected layer를 제거 self.fc = nn.Sequential( nn.Linear(in_features, 256), nn.ReLU(), nn.Linear(256, embedding_dim) ) def forward(self, x): h = self.base_model(x) z = self.fc(h) return F.normalize(z, dim=1) # 코사인 유사도 계산에 사용하기 위해 임베딩 벡터를 정규화
Python
복사

Objective

SimCLR의 loss 식을 따라 다음과 같이 구현한다.
여기서 out_1, out_2는 개별 이미지가 아니라 minibatch 전체이며 model을 통과해서 정규화된 데이터이다.
out_1과 out_2를 곱하면 positive 샘플의 유사도가 되고, 이 둘을 concatenate하여 만든 행렬에 대해 자기 자신과 행렬곱을 수행하면 대각 성분은 자기 자신과의 유사도가 되고, 나머지 부분은 negative 데이터와의 유사도가 된다.
SimCLR의 분모 부분에서 자기 자신의 유사도는 사용하지 않으므로 이 행렬의 대각 성분은 masking하여 최종 loss를 계산한다.
# 대조 손실 함수 정의 (코사인 유사도 사용) def SimCLR_loss(out_1, out_2, temperature): batch_size = out_1.shape[0] out = torch.cat([out_1, out_2], dim=0) # (2*batch_size, embedding_dim) # 코사인 유사도 행렬 계산 sim_matrix = torch.mm(out, out.t().contiguous()) / temperature sim_matrix = torch.exp(sim_matrix) # 대각 성분을 제거하기 위한 마스크 생성 mask = (torch.ones_like(sim_matrix) - torch.eye(2 * batch_size, device=sim_matrix.device)).bool() # 마스크 적용 sim_matrix = sim_matrix.masked_select(mask).view(2 * batch_size, -1) # 양성 샘플 유사도 계산 pos_sim = torch.sum(out_1 * out_2, dim=-1) / temperature pos_sim = torch.cat([pos_sim, pos_sim], dim=0) # 손실 계산 loss = pos_sim - torch.log(sim_matrix.sum(dim=-1)) return -loss.mean() # 손실 값을 최소화하기 위해 부호를 반전하고 평균을 계산
Python
복사

Train

데이터셋에 대해 데이터 증강을 수행한 후 모델을 통과시켜 embedding 결과를 얻고, 해당 결과에 대해 SimCLR 유사도와 손실을 계산한 다음 손실에 대해 역전파 하여 embedding을 업데이트한다.
# 데이터 증강 설정 train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8), transforms.RandomGrayscale(p=0.2), transforms.ToTensor() ]) # 데이터셋 로드 train_dataset = datasets.CIFAR10(root='./data', train=True, transform=train_transform, download=True) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4) # 모델 및 최적화 설정 base_model = models.resnet18(pretrained=False) model = SimCLRModel(base_model, base_model.fc.in_features, embedding_dim=128).to(device) optimizer = optim.Adam(model.parameters(), lr=3e-4) # 학습 num_epochs = 10 temperature = 0.5 for epoch in range(num_epochs): model.train() total_loss = 0 for (imgs, _) in train_loader: imgs = imgs.cuda() # 두 개의 변형된 버전을 생성 imgs_i = imgs imgs_j = imgs.clone() # 여기서 두 번째 변형된 이미지를 생성할 수 있습니다 out_1 = model(imgs_i) out_2 = model(imgs_j) loss = SimCLR_loss(out_1, out_2, temperature) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() avg_loss = total_loss / len(train_loader) print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}") print("Training completed.")
Python
복사

참고

Made with