lightly Quick Guide

Objective

run a sample code in lightly quickly

Download Dataset(CIFAR-10)

1. Register Kaggle
2. Download CIFAR-10 from Kaggle

Create Virtual Python Env

sudo apt install python3-pip
# sudo apt install python3-venv
python3.8 -m venv ~/venv/lightly
source ~/venv/lightly/bin/activate

Install lightly

pip install pip --upgrade
pip install torch
pip install torchvision
pip install pytorch-lightning
pip install lightly

Sample Code

• model: resnet18
• dataset: CIFAR-10
• self-supervised method: SimCLR

change path-to/cifar10

import torch
import torchvision

from lightly import loss
from lightly import transforms
from lightly.data import LightlyDataset
from lightly.models.modules import heads
from pytorch_lightning import LightningModule, Trainer


class SimCLR(LightningModule):
    def __init__(self):
        super().__init__()
        resnet = torchvision.models.resnet18()
        resnet.fc = torch.nn.Identity()
        self.backbone = resnet
        self.projection_head = heads.SimCLRProjectionHead(512, 512, 128)
        self.criterion = loss.NTXentLoss()

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

    def training_step(self, batch, batch_index):
        (view0, view1), _, _ = batch
        z0 = self.forward(view0)
        z1 = self.forward(view1)
        loss = self.criterion(z0, z1)
        return loss

    def configure_optimizers(self):
        optim = torch.optim.SGD(self.parameters(), lr=0.06)
        return optim

transform = transforms.SimCLRTransform(input_size=32, cj_prob=0.5)

dataset = LightlyDataset(input_dir="path-to/cifar10", transform=transform)

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=128,
    shuffle=True, 
)


model = SimCLR()
trainer = Trainer(max_epochs=10, devices=1, accelerator="gpu")
trainer.fit(model, dataloader)