def cli_main():
from ConSSL.callbacks.ssl_online import SSLOnlineEvaluator
from ConSSL.datamodules import CIFAR10DataModule, ImagenetDataModule, STL10DataModule
from ConSSL.models.self_supervised.swav.transforms import SwAVEvalDataTransform, SwAVTrainDataTransform
parser = ArgumentParser()
# model args
parser = SwAV.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
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 = 'sgd'
args.lars_wrapper = True
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,
)
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(): # pragma: no cover
from ConSSL.datamodules import ImagenetDataModule, STL10DataModule
pl.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.)
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.)
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 = pl.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 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 = 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.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.0
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 = 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.from_argparse_args(
args,
sync_batchnorm=True if args.gpus > 1 else False,
callbacks=[online_evaluator] if args.online_ft else None,
)
trainer.fit(model, datamodule=dm)