Python SwAV Examples

Example #1

File: backbones.py Project: alemelis/lightning-flash

 def load_swav_imagenet(
     path_or_url:
     str = f"{ROOT_S3_BUCKET}/swav/swav_imagenet/swav_imagenet.pth.tar"):
     swav = SwAV.load_from_checkpoint(path_or_url, strict=True)
     # remove the last two layers & turn it into a Sequential model
     backbone = nn.Sequential(*list(swav.model.children())[:-2])
     return backbone, 2048

Example #2

File: test_models.py Project: MihaiAnca13/pytorch-lightning-bolts

def test_swav(tmpdir, datadir, batch_size=2):
    # inputs, y = batch  (doesn't receive y for some reason)
    datamodule = CIFAR10DataModule(data_dir=datadir,
                                   batch_size=batch_size,
                                   num_workers=0)

    datamodule.train_transforms = SwAVTrainDataTransform(
        normalize=cifar10_normalization(),
        size_crops=[32, 16],
        nmb_crops=[2, 1],
        gaussian_blur=False)
    datamodule.val_transforms = SwAVEvalDataTransform(
        normalize=cifar10_normalization(),
        size_crops=[32, 16],
        nmb_crops=[2, 1],
        gaussian_blur=False)

    model = SwAV(arch='resnet18',
                 hidden_mlp=512,
                 gpus=0,
                 nodes=1,
                 num_samples=datamodule.num_samples,
                 batch_size=batch_size,
                 nmb_crops=[2, 1],
                 sinkhorn_iterations=1,
                 nmb_prototypes=2,
                 queue_length=0,
                 maxpool1=False,
                 first_conv=False,
                 dataset='cifar10')

    trainer = Trainer(gpus=0, fast_dev_run=True, default_root_dir=tmpdir)

    trainer.fit(model, datamodule=datamodule)

Example #3

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()

Example #4

File: test_models.py Project: shamazharikh/pytorch-lightning-bolts

def test_swav(tmpdir, datadir):
    seed_everything()

    batch_size = 2

    # inputs, y = batch  (doesn't receive y for some reason)
    datamodule = CIFAR10DataModule(
        data_dir=datadir,
        batch_size=batch_size,
        num_workers=0
    )

    datamodule.train_transforms = SwAVTrainDataTransform(
        normalize=cifar10_normalization(),
        size_crops=[32, 16],
        nmb_crops=[2, 1],
        gaussian_blur=False
    )
    datamodule.val_transforms = SwAVEvalDataTransform(
        normalize=cifar10_normalization(),
        size_crops=[32, 16],
        nmb_crops=[2, 1],
        gaussian_blur=False
    )

    model = SwAV(
        arch='resnet18',
        hidden_mlp=512,
        gpus=0,
        nodes=1,
        num_samples=datamodule.num_samples,
        batch_size=batch_size,
        nmb_crops=[2, 1],
        sinkhorn_iterations=1,
        nmb_prototypes=2,
        maxpool1=False,
        first_conv=False,
        dataset='cifar10'
    )

    trainer = pl.Trainer(
        gpus=0, fast_dev_run=True, max_epochs=1, default_root_dir=tmpdir, max_steps=3
    )

    trainer.fit(model, datamodule)
    loss = trainer.progress_bar_dict['loss']

    assert float(loss) > 0

Example #5

File: model_map.py Project: peteryonka/lightning-flash

def load_swav_imagenet(
    path_or_url:
    str = f"{ROOT_S3_BUCKET}/swav/swav_imagenet/swav_imagenet.pth.tar"):
    swav = SwAV.load_from_checkpoint(path_or_url, strict=True)
    model_config = {'model': swav.model, 'num_features': 3000}
    return model_config