train_encoder.py

import pickle

import torch
import torchvision
from torch import nn

from lightly.loss import NTXentLoss
from lightly.models.modules import SimCLRProjectionHead
from lightly.transforms.simclr_transform import SimCLRTransform


class SimCLR(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.projection_head = SimCLRProjectionHead(512, 512, 128)

    def forward(self, x):
        x = self.backbone(x).flatten(start_dim=1)
        z = self.projection_head(x)
        return z



if __name__ == "__main__":
        resnet = torchvision.models.resnet18()
        resnet.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        backbone = nn.Sequential(*list(resnet.children())[:-1])
        model = SimCLR(backbone)
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)
        
        transform = SimCLRTransform(input_size=28, gaussian_blur=0.0, normalize={"mean": [0.5], "std": [0.5]})
        dataset = torchvision.datasets.MNIST(
            root="data", train=True, download=True, transform=transform
        )
        
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=256,
            shuffle=True,
            drop_last=True,
            num_workers=8,
        )
        
        criterion = NTXentLoss()
        optimizer = torch.optim.SGD(model.parameters(), lr=0.06)
        
        print("Starting Training")
        loss_logs = []
        for epoch in range(200):
            total_loss = 0
            for batch in dataloader:
                x0, x1 = batch[0]
                x0 = x0.to(device)
                x1 = x1.to(device)
                z0 = model(x0)
                z1 = model(x1)
                loss = criterion(z0, z1)
                total_loss += loss.detach()
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
            avg_loss = total_loss / len(dataloader)
            print(f"epoch: {epoch:>02}, loss: {avg_loss:.5f}")
            loss_logs.append(avg_loss.item())
        
        
        # save logs and model
        with open('logs/sim_clr_loss_log.pkl', 'wb') as f:
            pickle.dump(loss_logs, f)
        
        torch.save(model.backbone.state_dict(), "models/encoder_mnist.pth")