Skip to content

Instantly share code, notes, and snippets.

@williamFalcon

williamFalcon/simclr.py

Created July 22, 2020 15:13
Show Gist options
  • Save williamFalcon/b283ce8842a63aead369f9692194f9b9 to your computer and use it in GitHub Desktop.
Save williamFalcon/b283ce8842a63aead369f9692194f9b9 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment