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)
class VAE(pl.LightningModule): def __init__(self, hidden_dim: int = 128, latent_dim: int = 32, input_channels: int = 3, input_width: int = 224, input_height: int = 224, batch_size: int = 32, learning_rate: float = 0.001, data_dir: str = '.', datamodule: pl.LightningDataModule = None, num_workers: int = 8, pretrained: str = None, **kwargs): """ Standard VAE with Gaussian Prior and approx posterior. Model is available pretrained on different datasets: Example:: # not pretrained vae = VAE() # pretrained on imagenet vae = VAE(pretrained='imagenet') # pretrained on cifar10 vae = VAE(pretrained='cifar10') Args: hidden_dim: encoder and decoder hidden dims latent_dim: latenet code dim input_channels: num of channels of the input image. input_width: image input width input_height: image input height batch_size: the batch size learning_rate" the learning rate data_dir: the directory to store data datamodule: The Lightning DataModule pretrained: Load weights pretrained on a dataset """ super().__init__() self.save_hyperparameters() self.datamodule = datamodule self.__set_pretrained_dims(pretrained) # use mnist as the default module self._set_default_datamodule(datamodule) # init actual model self.__init_system() if pretrained: self.load_pretrained(pretrained) def __init_system(self): self.encoder = self.init_encoder() self.decoder = self.init_decoder() def __set_pretrained_dims(self, pretrained): if pretrained == 'imagenet2012': self.datamodule = ImagenetDataModule( data_dir=self.hparams.data_dir) (self.hparams.input_channels, self.hparams.input_height, self.hparams.input_width) = self.datamodule.size() def _set_default_datamodule(self, datamodule): # link default data if datamodule is None: datamodule = MNISTDataModule(data_dir=self.hparams.data_dir, num_workers=self.hparams.num_workers, normalize=False) self.datamodule = datamodule self.img_dim = self.datamodule.size() (self.hparams.input_channels, self.hparams.input_height, self.hparams.input_width) = self.img_dim def load_pretrained(self, pretrained): available_weights = {'imagenet2012'} if pretrained in available_weights: weights_name = f'vae-{pretrained}' load_pretrained(self, weights_name) def init_encoder(self): encoder = Encoder(self.hparams.hidden_dim, self.hparams.latent_dim, self.hparams.input_channels, self.hparams.input_width, self.hparams.input_height) return encoder def init_decoder(self): decoder = Decoder(self.hparams.hidden_dim, self.hparams.latent_dim, self.hparams.input_width, self.hparams.input_height, self.hparams.input_channels) return decoder def get_prior(self, z_mu, z_std): # Prior ~ Normal(0,1) P = distributions.normal.Normal(loc=torch.zeros_like(z_mu), scale=torch.ones_like(z_std)) return P def get_approx_posterior(self, z_mu, z_std): # Approx Posterior ~ Normal(mu, sigma) Q = distributions.normal.Normal(loc=z_mu, scale=z_std) return Q def elbo_loss(self, x, P, Q, num_samples): z = Q.rsample() # ---------------------- # KL divergence loss (using monte carlo sampling) # ---------------------- log_qz = Q.log_prob(z) log_pz = P.log_prob(z) # (batch, num_samples, z_dim) -> (batch, num_samples) kl_div = (log_qz - log_pz).sum(dim=2) # we used monte carlo sampling to estimate. average across samples # kl_div = kl_div.mean(-1) # ---------------------- # Reconstruction loss # ---------------------- z = z.view(-1, z.size(-1)).contiguous() pxz = self.decoder(z) pxz = pxz.view(-1, num_samples, pxz.size(-1)) x = shaping.tile(x.unsqueeze(1), 1, num_samples) pxz = torch.sigmoid(pxz) recon_loss = F.binary_cross_entropy(pxz, x, reduction='none') # sum across dimensions because sum of log probabilities of iid univariate gaussians is the same as # multivariate gaussian recon_loss = recon_loss.sum(dim=-1) # we used monte carlo sampling to estimate. average across samples # recon_loss = recon_loss.mean(-1) # ELBO = reconstruction + KL loss = recon_loss + kl_div # average over batch loss = loss.mean() recon_loss = recon_loss.mean() kl_div = kl_div.mean() return loss, recon_loss, kl_div, pxz def forward(self, z): return self.decoder(z) def _run_step(self, batch): x, _ = batch z_mu, z_log_var = self.encoder(x) # we're estimating the KL divergence using sampling num_samples = 32 # expand dims to sample all at once # (batch, z_dim) -> (batch, num_samples, z_dim) z_mu = z_mu.unsqueeze(1) z_mu = shaping.tile(z_mu, 1, num_samples) # (batch, z_dim) -> (batch, num_samples, z_dim) z_log_var = z_log_var.unsqueeze(1) z_log_var = shaping.tile(z_log_var, 1, num_samples) # convert to std z_std = torch.exp(z_log_var / 2) P = self.get_prior(z_mu, z_std) Q = self.get_approx_posterior(z_mu, z_std) x = x.view(x.size(0), -1) loss, recon_loss, kl_div, pxz = self.elbo_loss(x, P, Q, num_samples) return loss, recon_loss, kl_div, pxz def training_step(self, batch, batch_idx): loss, recon_loss, kl_div, pxz = self._run_step(batch) result = pl.TrainResult(loss) result.log_dict({ 'train_elbo_loss': loss, 'train_recon_loss': recon_loss, 'train_kl_loss': kl_div }) return result def validation_step(self, batch, batch_idx): loss, recon_loss, kl_div, pxz = self._run_step(batch) result = pl.EvalResult(loss, checkpoint_on=loss) result.log_dict({ 'val_loss': loss, 'val_recon_loss': recon_loss, 'val_kl_div': kl_div, }) return result def test_step(self, batch, batch_idx): loss, recon_loss, kl_div, pxz = self._run_step(batch) result = pl.EvalResult(loss) result.log_dict({ 'test_loss': loss, 'test_recon_loss': recon_loss, 'test_kl_div': kl_div, }) return result def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate) @staticmethod def add_model_specific_args(parent_parser): parser = ArgumentParser(parents=[parent_parser], add_help=False) parser.add_argument( '--hidden_dim', type=int, default=128, help= 'itermediate layers dimension before embedding for default encoder/decoder' ) parser.add_argument('--latent_dim', type=int, default=4, help='dimension of latent variables z') parser.add_argument( '--input_width', type=int, default=224, help='input width (used Imagenet downsampled size)') parser.add_argument( '--input_height', type=int, default=224, help='input width (used Imagenet downsampled size)') parser.add_argument('--input_channels', type=int, default=3, help='number of input channels') parser.add_argument('--batch_size', type=int, default=64) parser.add_argument('--pretrained', type=str, default=None) parser.add_argument('--data_dir', type=str, default=os.getcwd()) parser.add_argument('--num_workers', type=int, default=8, help="num dataloader workers") parser.add_argument('--learning_rate', type=float, default=1e-3) return parser
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.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)