def swav_example(): from pl_bolts.models.self_supervised import SwAV from pl_bolts.datamodules import STL10DataModule from pl_bolts.models.self_supervised.swav.transforms import SwAVTrainDataTransform, SwAVEvalDataTransform from pl_bolts.transforms.dataset_normalizations import stl10_normalization batch_size = 128 # Data module. dm = STL10DataModule(data_dir=".", num_workers=16, batch_size=batch_size) dm.train_dataloader = dm.train_dataloader_mixed dm.val_dataloader = dm.val_dataloader_mixed dm.train_transforms = SwAVTrainDataTransform(normalize=stl10_normalization()) dm.val_transforms = SwAVEvalDataTransform(normalize=stl10_normalization()) # Model. model = SwAV( gpus=1, num_samples=dm.num_unlabeled_samples, dataset="stl10", batch_size=batch_size ) # Fit. trainer = pl.Trainer(gpus=2, accelerator="ddp", precision=16) trainer.fit(model, datamodule=dm) #-------------------- # ImageNet pretrained model: weight_path = "https://pl-bolts-weights.s3.us-east-2.amazonaws.com/swav/bolts_swav_imagenet/swav_imagenet.ckpt" #weight_path = "https://pl-bolts-weights.s3.us-east-2.amazonaws.com/swav/swav_imagenet/swav_imagenet.pth.tar" # STL-10 pretrained model: #weight_path = "https://pl-bolts-weights.s3.us-east-2.amazonaws.com/swav/checkpoints/swav_stl10.pth.tar" swav = SwAV.load_from_checkpoint(weight_path, strict=True) swav.freeze()
def cli_main(): # pragma: no cover from pl_bolts.datamodules import ImagenetDataModule, STL10DataModule seed_everything(1234) parser = ArgumentParser() parser.add_argument("--dataset", type=str, help="stl10, imagenet", default="stl10") parser.add_argument("--ckpt_path", type=str, help="path to ckpt") parser.add_argument("--data_dir", type=str, help="path to dataset", default=os.getcwd()) parser.add_argument("--batch_size", default=64, type=int, help="batch size per gpu") parser.add_argument("--num_workers", default=8, type=int, help="num of workers per GPU") parser.add_argument("--gpus", default=4, type=int, help="number of GPUs") parser.add_argument("--num_epochs", default=100, type=int, help="number of epochs") # fine-tuner params parser.add_argument("--in_features", type=int, default=2048) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--learning_rate", type=float, default=0.3) parser.add_argument("--weight_decay", type=float, default=1e-6) parser.add_argument("--nesterov", type=bool, default=False) parser.add_argument("--scheduler_type", type=str, default="cosine") parser.add_argument("--gamma", type=float, default=0.1) parser.add_argument("--final_lr", type=float, default=0.0) args = parser.parse_args() if args.dataset == "stl10": dm = STL10DataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) dm.train_dataloader = dm.train_dataloader_labeled dm.val_dataloader = dm.val_dataloader_labeled args.num_samples = 0 dm.train_transforms = SwAVFinetuneTransform( normalize=stl10_normalization(), input_height=dm.size()[-1], eval_transform=False) dm.val_transforms = SwAVFinetuneTransform( normalize=stl10_normalization(), input_height=dm.size()[-1], eval_transform=True) dm.test_transforms = SwAVFinetuneTransform( normalize=stl10_normalization(), input_height=dm.size()[-1], eval_transform=True) args.maxpool1 = False args.first_conv = True elif args.dataset == "imagenet": dm = ImagenetDataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) dm.train_transforms = SwAVFinetuneTransform( normalize=imagenet_normalization(), input_height=dm.size()[-1], eval_transform=False) dm.val_transforms = SwAVFinetuneTransform( normalize=imagenet_normalization(), input_height=dm.size()[-1], eval_transform=True) dm.test_transforms = SwAVFinetuneTransform( normalize=imagenet_normalization(), input_height=dm.size()[-1], eval_transform=True) args.num_samples = 1 args.maxpool1 = True args.first_conv = True else: raise NotImplementedError( "other datasets have not been implemented till now") backbone = SwAV( gpus=args.gpus, nodes=1, num_samples=args.num_samples, batch_size=args.batch_size, maxpool1=args.maxpool1, first_conv=args.first_conv, dataset=args.dataset, ).load_from_checkpoint(args.ckpt_path, strict=False) tuner = SSLFineTuner( backbone, in_features=args.in_features, num_classes=dm.num_classes, epochs=args.num_epochs, hidden_dim=None, dropout=args.dropout, learning_rate=args.learning_rate, weight_decay=args.weight_decay, nesterov=args.nesterov, scheduler_type=args.scheduler_type, gamma=args.gamma, final_lr=args.final_lr, ) trainer = Trainer( gpus=args.gpus, num_nodes=1, precision=16, max_epochs=args.num_epochs, distributed_backend="ddp", sync_batchnorm=True if args.gpus > 1 else False, ) trainer.fit(tuner, dm) trainer.test(datamodule=dm)
def cli_main(): from pl_bolts.callbacks.ssl_online import SSLOnlineEvaluator from pl_bolts.datamodules import CIFAR10DataModule, ImagenetDataModule, STL10DataModule from pl_bolts.models.self_supervised.swav.transforms import SwAVEvalDataTransform, SwAVTrainDataTransform parser = ArgumentParser() # model args parser = SwAV.add_model_specific_args(parser) args = parser.parse_args() if args.dataset == 'stl10': dm = STL10DataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) dm.train_dataloader = dm.train_dataloader_mixed dm.val_dataloader = dm.val_dataloader_mixed args.num_samples = dm.num_unlabeled_samples args.maxpool1 = False normalization = stl10_normalization() elif args.dataset == 'cifar10': args.batch_size = 2 args.num_workers = 0 dm = CIFAR10DataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) args.num_samples = dm.num_samples args.maxpool1 = False args.first_conv = False normalization = cifar10_normalization() # cifar10 specific params args.size_crops = [32, 16] args.nmb_crops = [2, 1] args.gaussian_blur = False elif args.dataset == 'imagenet': args.maxpool1 = True args.first_conv = True normalization = imagenet_normalization() args.size_crops = [224, 96] args.nmb_crops = [2, 6] args.min_scale_crops = [0.14, 0.05] args.max_scale_crops = [1., 0.14] args.gaussian_blur = True args.jitter_strength = 1. args.batch_size = 64 args.num_nodes = 8 args.gpus = 8 # per-node args.max_epochs = 800 args.optimizer = 'lars' args.learning_rate = 4.8 args.final_lr = 0.0048 args.start_lr = 0.3 args.nmb_prototypes = 3000 args.online_ft = True dm = ImagenetDataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) args.num_samples = dm.num_samples args.input_height = dm.size()[-1] else: raise NotImplementedError( "other datasets have not been implemented till now") dm.train_transforms = SwAVTrainDataTransform( normalize=normalization, size_crops=args.size_crops, nmb_crops=args.nmb_crops, min_scale_crops=args.min_scale_crops, max_scale_crops=args.max_scale_crops, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength) dm.val_transforms = SwAVEvalDataTransform( normalize=normalization, size_crops=args.size_crops, nmb_crops=args.nmb_crops, min_scale_crops=args.min_scale_crops, max_scale_crops=args.max_scale_crops, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength) # swav model init model = SwAV(**args.__dict__) online_evaluator = None if args.online_ft: # online eval online_evaluator = SSLOnlineEvaluator( drop_p=0., hidden_dim=None, z_dim=args.hidden_mlp, num_classes=dm.num_classes, dataset=args.dataset, ) lr_monitor = LearningRateMonitor(logging_interval="step") model_checkpoint = ModelCheckpoint(save_last=True, save_top_k=1, monitor='val_loss') callbacks = [model_checkpoint, online_evaluator ] if args.online_ft else [model_checkpoint] callbacks.append(lr_monitor) trainer = Trainer( max_epochs=args.max_epochs, max_steps=None if args.max_steps == -1 else args.max_steps, gpus=args.gpus, num_nodes=args.num_nodes, distributed_backend='ddp' if args.gpus > 1 else None, sync_batchnorm=True if args.gpus > 1 else False, precision=32 if args.fp32 else 16, callbacks=callbacks, fast_dev_run=args.fast_dev_run) trainer.fit(model, datamodule=dm)
def cli_main(): from pl_bolts.callbacks.ssl_online import SSLOnlineEvaluator from pl_bolts.datamodules import CIFAR10DataModule, ImagenetDataModule, STL10DataModule from pl_bolts.models.self_supervised.simclr.transforms import SimCLREvalDataTransform, SimCLRTrainDataTransform parser = ArgumentParser() # model args parser = SimCLR.add_model_specific_args(parser) parser = pl.Trainer.add_argparse_args(parser) args = parser.parse_args() if args.dataset == 'stl10': dm = STL10DataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) dm.train_dataloader = dm.train_dataloader_mixed dm.val_dataloader = dm.val_dataloader_mixed args.num_samples = dm.num_unlabeled_samples args.maxpool1 = False args.first_conv = True args.input_height = dm.size()[-1] normalization = stl10_normalization() args.gaussian_blur = True args.jitter_strength = 1. elif args.dataset == 'cifar10': val_split = 5000 if args.num_nodes * args.gpus * args.batch_size > val_split: val_split = args.num_nodes * args.gpus * args.batch_size dm = CIFAR10DataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers, val_split=val_split) args.num_samples = dm.num_samples args.maxpool1 = False args.first_conv = False args.input_height = dm.size()[-1] args.temperature = 0.5 normalization = cifar10_normalization() args.gaussian_blur = False args.jitter_strength = 0.5 elif args.dataset == 'imagenet': args.maxpool1 = True args.first_conv = True normalization = imagenet_normalization() args.gaussian_blur = True args.jitter_strength = 1. args.batch_size = 64 args.num_nodes = 8 args.gpus = 8 # per-node args.max_epochs = 800 args.optimizer = 'sgd' args.lars_wrapper = True args.learning_rate = 4.8 args.final_lr = 0.0048 args.start_lr = 0.3 args.online_ft = True dm = ImagenetDataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) args.num_samples = dm.num_samples args.input_height = dm.size()[-1] else: raise NotImplementedError( "other datasets have not been implemented till now") dm.train_transforms = SimCLRTrainDataTransform( input_height=args.input_height, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength, normalize=normalization, ) dm.val_transforms = SimCLREvalDataTransform( input_height=args.input_height, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength, normalize=normalization, ) model = SimCLR(**args.__dict__) online_evaluator = None if args.online_ft: # online eval online_evaluator = SSLOnlineEvaluator(drop_p=0., hidden_dim=None, z_dim=args.hidden_mlp, num_classes=dm.num_classes, dataset=args.dataset) model_checkpoint = ModelCheckpoint(save_last=True, save_top_k=1, monitor='val_loss') callbacks = [model_checkpoint, online_evaluator ] if args.online_ft else [model_checkpoint] trainer = pl.Trainer.from_argparse_args( args, sync_batchnorm=True if args.gpus > 1 else False, callbacks=callbacks, ) trainer.fit(model, datamodule=dm)
def cli_main(): from pl_bolts.callbacks.ssl_online import SSLOnlineEvaluator from pl_bolts.models.self_supervised.swav.transforms import SwAVTrainDataTransform, SwAVEvalDataTransform from pl_bolts.datamodules import STL10DataModule, CIFAR10DataModule parser = ArgumentParser() # model args parser = SwAV.add_model_specific_args(parser) args = parser.parse_args() if args.dataset == 'stl10': dm = STL10DataModule(data_dir=args.data_path, batch_size=args.batch_size, num_workers=args.num_workers) dm.train_dataloader = dm.train_dataloader_mixed dm.val_dataloader = dm.val_dataloader_mixed args.num_samples = dm.num_unlabeled_samples args.maxpool1 = False normalization = stl10_normalization() elif args.dataset == 'cifar10': args.batch_size = 2 args.num_workers = 0 dm = CIFAR10DataModule(data_dir=args.data_path, batch_size=args.batch_size, num_workers=args.num_workers) args.num_samples = dm.num_samples args.maxpool1 = False args.first_conv = False normalization = cifar10_normalization() # cifar10 specific params args.size_crops = [32, 16] args.nmb_crops = [2, 1] args.gaussian_blur = False else: raise NotImplementedError( "other datasets have not been implemented till now") dm.train_transforms = SwAVTrainDataTransform( normalize=stl10_normalization(), size_crops=args.size_crops, nmb_crops=args.nmb_crops, min_scale_crops=args.min_scale_crops, max_scale_crops=args.max_scale_crops, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength) dm.val_transforms = SwAVEvalDataTransform( normalize=stl10_normalization(), size_crops=args.size_crops, nmb_crops=args.nmb_crops, min_scale_crops=args.min_scale_crops, max_scale_crops=args.max_scale_crops, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength) # swav model init model = SwAV(**args.__dict__) online_evaluator = None if args.online_ft: # online eval online_evaluator = SSLOnlineEvaluator(drop_p=0., hidden_dim=None, z_dim=args.hidden_mlp, num_classes=dm.num_classes, dataset=args.dataset) trainer = pl.Trainer( max_epochs=args.max_epochs, max_steps=None if args.max_steps == -1 else args.max_steps, gpus=args.gpus, distributed_backend='ddp' if args.gpus > 1 else None, sync_batchnorm=True if args.gpus > 1 else False, precision=32 if args.fp32 else 16, callbacks=[online_evaluator] if args.online_ft else None, fast_dev_run=args.fast_dev_run) trainer.fit(model, dm)
import pytorch_lightning as pl from swav_module import SwAV from pl_bolts.datamodules import STL10DataModule from pl_bolts.models.self_supervised.swav.transforms import ( SwAVTrainDataTransform, SwAVEvalDataTransform ) from pl_bolts.transforms.dataset_normalizations import stl10_normalization # data batch_size = 128 dm = STL10DataModule(data_dir='.', batch_size=batch_size, num_workers=8) dm.train_dataloader = dm.train_dataloader_mixed dm.val_dataloader = dm.val_dataloader_mixed dm.train_transforms = SwAVTrainDataTransform( normalize=stl10_normalization() ) dm.val_transforms = SwAVEvalDataTransform( normalize=stl10_normalization() ) # model model = SwAV( lars_wrapper=True, online_ft=True, gpus=1, learning_rate=1e-3, num_samples=dm.num_unlabeled_samples, gaussian_blur=True, queue_length=0,
def cli_main(): from pl_bolts.callbacks.ssl_online import SSLOnlineEvaluator from pl_bolts.datamodules import CIFAR10DataModule, ImagenetDataModule, STL10DataModule from pl_bolts.models.self_supervised.simclr import SimCLREvalDataTransform, SimCLRTrainDataTransform seed_everything(1234) parser = ArgumentParser() # trainer args parser = pl.Trainer.add_argparse_args(parser) # model args parser = SimSiam.add_model_specific_args(parser) args = parser.parse_args() # pick data dm = None # init datamodule if args.dataset == "stl10": dm = STL10DataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) dm.train_dataloader = dm.train_dataloader_mixed dm.val_dataloader = dm.val_dataloader_mixed args.num_samples = dm.num_unlabeled_samples args.maxpool1 = False args.first_conv = True args.input_height = dm.size()[-1] normalization = stl10_normalization() args.gaussian_blur = True args.jitter_strength = 1.0 elif args.dataset == "cifar10": val_split = 5000 if args.nodes * args.gpus * args.batch_size > val_split: val_split = args.nodes * args.gpus * args.batch_size dm = CIFAR10DataModule( data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers, val_split=val_split, ) args.num_samples = dm.num_samples args.maxpool1 = False args.first_conv = False args.input_height = dm.size()[-1] args.temperature = 0.5 normalization = cifar10_normalization() args.gaussian_blur = False args.jitter_strength = 0.5 elif args.dataset == "cifar100": val_split = 5000 if args.nodes * args.gpus * args.batch_size > val_split: val_split = args.nodes * args.gpus * args.batch_size dm = CIFAR100DataModule( data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers, val_split=val_split, ) args.num_samples = dm.num_samples args.maxpool1 = False args.first_conv = False args.input_height = dm.size()[-1] args.temperature = 0.5 # ((0.5071, 0.4866, 0.4409), (0.2009, 0.1984, 0.2023)) normalization = transforms.Normalize( mean=(0.5071, 0.4866, 0.4409), std=(0.2009, 0.1984, 0.2023), ) args.gaussian_blur = False args.jitter_strength = 0.5 elif args.dataset == "imagenet": args.maxpool1 = True args.first_conv = True normalization = imagenet_normalization() args.gaussian_blur = True args.jitter_strength = 1.0 args.batch_size = 64 args.nodes = 8 args.gpus = 8 # per-node args.max_epochs = 800 args.optimizer = "sgd" args.lars_wrapper = True args.learning_rate = 4.8 args.final_lr = 0.0048 args.start_lr = 0.3 args.online_ft = True dm = ImagenetDataModule(data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers) args.num_samples = dm.num_samples args.input_height = dm.size()[-1] else: raise NotImplementedError("other datasets have not been implemented till now") dm.train_transforms = SimCLRTrainDataTransform( input_height=args.input_height, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength, normalize=normalization, ) dm.val_transforms = SimCLREvalDataTransform( input_height=args.input_height, gaussian_blur=args.gaussian_blur, jitter_strength=args.jitter_strength, normalize=normalization, ) model = SimSiam(**args.__dict__) # finetune in real-time online_evaluator = None if args.online_ft: # online eval online_evaluator = SSLOnlineEvaluator( drop_p=0.0, hidden_dim=None, z_dim=args.hidden_mlp, num_classes=dm.num_classes, dataset=args.dataset, ) trainer = pl.Trainer( max_epochs=args.max_epochs, max_steps=None if args.max_steps == -1 else args.max_steps, gpus=args.gpus, num_nodes=args.nodes, distributed_backend="ddp" if args.gpus > 1 else None, sync_batchnorm=True if args.gpus > 1 else False, precision=32 if args.fp32 else 16, callbacks=[online_evaluator] if args.online_ft else None, fast_dev_run=args.fast_dev_run, ) trainer.fit(model, dm)