def cli_main():
from ConSSL.callbacks.ssl_online import SSLOnlineEvaluator
from ConSSL.datamodules import CIFAR10DataModule, ImagenetDataModule, STL10DataModule
from ConSSL.models.self_supervised.simclr import SimCLREvalDataTransform, SimCLRTrainDataTransform
seed_everything(1234)
parser = ArgumentParser()
# trainer args
parser = pl.Trainer.add_argparse_args(parser)
# model args
parser = BYOL.add_model_specific_args(parser)
args = parser.parse_args()
# pick data
dm = None
# init default datamodule
if args.dataset == 'cifar10':
dm = CIFAR10DataModule.from_argparse_args(args)
dm.train_transforms = SimCLRTrainDataTransform(32)
dm.val_transforms = SimCLREvalDataTransform(32)
args.num_classes = dm.num_classes
elif args.dataset == 'stl10':
dm = STL10DataModule.from_argparse_args(args)
dm.train_dataloader = dm.train_dataloader_mixed
dm.val_dataloader = dm.val_dataloader_mixed
(c, h, w) = dm.size()
dm.train_transforms = SimCLRTrainDataTransform(h)
dm.val_transforms = SimCLREvalDataTransform(h)
args.num_classes = dm.num_classes
elif args.dataset == 'imagenet2012':
dm = ImagenetDataModule.from_argparse_args(args, image_size=196)
(c, h, w) = dm.size()
dm.train_transforms = SimCLRTrainDataTransform(h)
dm.val_transforms = SimCLREvalDataTransform(h)
args.num_classes = dm.num_classes
model = BYOL(**args.__dict__)
# finetune in real-time
online_eval = SSLOnlineEvaluator(dataset=args.dataset,
z_dim=2048,
num_classes=dm.num_classes)
trainer = pl.Trainer.from_argparse_args(args,
max_steps=300000,
callbacks=[online_eval])
trainer.fit(model, datamodule=dm)
def cli_main():
from ConSSL.datamodules import FashionMNISTDataModule, ImagenetDataModule
parser = ArgumentParser()
# trainer args
parser = pl.Trainer.add_argparse_args(parser)
# model args
parser = ImageGPT.add_model_specific_args(parser)
args = parser.parse_args()
if args.dataset == "fashion_mnist":
datamodule = FashionMNISTDataModule.from_argparse_args(args)
elif args.dataset == "imagenet128":
datamodule = ImagenetDataModule.from_argparse_args(args)
model = ImageGPT(**args.__dict__)
trainer = pl.Trainer.from_argparse_args(args)
trainer.fit(model, datamodule=datamodule)