simclr.py

import pytorch_lightning as pl
from pl_bolts.models.self_supervised import SimCLR
from pl_bolts.datamodules import ImagenetDataModule
from pl_bolts.models.self_supervised.simclr.transforms import SimCLRTrainDataTransform, SimCLREvalDataTransform

# data
datamodule = ImagenetDataModule(image_size=196)

# transforms
(c, h, w) = datamodule.size()
datamodule.train_transforms = SimCLRTrainDataTransform(h)
datamodule.val_transforms = SimCLREvalDataTransform(h)

# model
model = SimCLR(datamodule=datamodule)

# expensive!
trainer = pl.Trainer(tpu_cores=128)
trainer.fit(model