diff options
| -rw-r--r-- | readme.md | 2 | ||||
| -rw-r--r-- | simclr/evaluate.py | 268 | ||||
| -rw-r--r-- | simclr/models.py | 42 |
3 files changed, 295 insertions, 17 deletions
@@ -16,7 +16,7 @@ - [x] Command line parameter parser - [x] YAML parameter parser - [x] PyTorch workflow encapsulation - - [ ] Dedicated evaluation script + - [x] Dedicated evaluation script - [ ] Detailed readme file - [ ] Data Parallel - [ ] Distributed Data Parallel diff --git a/simclr/evaluate.py b/simclr/evaluate.py new file mode 100644 index 0000000..c029ffe --- /dev/null +++ b/simclr/evaluate.py @@ -0,0 +1,268 @@ +import argparse +import os.path +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Callable + +import sys +import torch +import yaml +from torch.utils.data import Dataset +from torchvision.datasets import CIFAR10, CIFAR100, ImageNet +from torchvision.transforms import transforms + +path = str(Path(Path(__file__).parent.absolute()).parent.absolute()) +sys.path.insert(0, path) + +from libs.optimizers import LARS +from libs.logging import Loggers, BaseBatchLogRecord, BaseEpochLogRecord +from libs.utils import BaseConfig +from simclr.main import SimCLRTrainer, SimCLRConfig +from simclr.models import CIFARSimCLRResNet50, ImageNetSimCLRResNet50 + + +def parse_args_and_config(): + parser = argparse.ArgumentParser( + description='Contrastive baseline SimCLR (evaluation)', + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument('--codename', default='cifar10-simclr-linear', + type=str, help="Model descriptor") + parser.add_argument('--log-dir', default='logs', type=str, + help="Path to log directory") + parser.add_argument('--checkpoint-dir', default='checkpoints', type=str, + help="Path to checkpoints directory") + parser.add_argument('--seed', default=None, type=int, + help='Random seed for reproducibility') + parser.add_argument('--num-iters', default=50, type=int, + help='Number of iters') + parser.add_argument('--config', type=argparse.FileType(mode='r'), + help='Path to config file (optional)') + + # TODO: Add model hyperparams dataclass + parser.add_argument('--hid-dim', default=2048, type=int, + help='Number of dimension of embedding') + parser.add_argument('--out-dim', default=128, type=int, + help='Number of dimension after projection') + parser.add_argument('--pretrained-checkpoint', type=str, + help='Pretrained checkpoint location') + parser.add_argument('--finetune', default=False, + action=argparse.BooleanOptionalAction, + help='Finetune backbone or linear head only') + + dataset_group = parser.add_argument_group('Dataset parameters') + dataset_group.add_argument('--dataset-dir', default='dataset', type=str, + help="Path to dataset directory") + dataset_group.add_argument('--dataset', default='cifar10', type=str, + choices=('cifar10, cifar100', 'imagenet'), + help="Name of dataset") + dataset_group.add_argument('--train-size', default=32, type=int, + help='Resize during training') + dataset_group.add_argument('--test-size', default=32, type=int, + help='Resize during testing') + dataset_group.add_argument('--test-crop-size', default=32, type=int, + help='Center crop size during testing') + dataset_group.add_argument('--hflip-prob', default=0.5, type=float, + help='Random horizontal flip probability') + + dataloader_group = parser.add_argument_group('Dataloader parameters') + dataloader_group.add_argument('--batch-size', default=256, type=int, + help='Batch size') + dataloader_group.add_argument('--num-workers', default=2, type=int, + help='Number of dataloader processes') + + optim_group = parser.add_argument_group('Optimizer parameters') + optim_group.add_argument('--optim', default='sgd', type=str, + choices=('adam', 'sgd', 'lars'), + help="Name of optimizer") + optim_group.add_argument('--lr', default=1e-3, type=float, + help='Learning rate') + optim_group.add_argument('--betas', nargs=2, default=(0.9, 0.999), type=float, + help='Adam betas', metavar=('beta1', 'beta2')) + optim_group.add_argument('--momentum', default=0.9, type=float, + help='SDG momentum') + optim_group.add_argument('--weight-decay', default=0., type=float, + help='Weight decay (l2 regularization)') + + sched_group = parser.add_argument_group('Optimizer parameters') + sched_group.add_argument('--sched', default=None, type=str, + help="Name of scheduler") + sched_group.add_argument('--warmup-iters', default=5, type=int, + help='Epochs for warmup (`warmup-anneal` scheduler only)') + + args = parser.parse_args() + if args.config: + config = yaml.safe_load(args.config) + args.__dict__ |= { + k: tuple(v) if isinstance(v, list) else v + for k, v in config.items() + } + args.checkpoint_dir = os.path.join(args.checkpoint_dir, args.codename) + args.log_dir = os.path.join(args.log_dir, args.codename) + + return args + + +@dataclass +class SimCLREvalConfig(SimCLRConfig): + @dataclass + class DatasetConfig(BaseConfig.DatasetConfig): + dataset_dir: str + train_size: int + test_size: int + test_crop_size: int + hflip_prob: float + + +class SimCLREvalTrainer(SimCLRTrainer): + def __init__(self, pretrained_checkpoint, finetune, **kwargs): + self.pretrained_checkpoint = pretrained_checkpoint + self.finetune = finetune + super(SimCLREvalTrainer, self).__init__(**kwargs) + + @dataclass + class BatchLogRecord(BaseBatchLogRecord): + lr: float + train_loss: float + + @dataclass + class EpochLogRecord(BaseEpochLogRecord): + eval_loss: float + eval_accuracy: float + + @staticmethod + def _prepare_dataset(dataset_config: SimCLREvalConfig.DatasetConfig) -> tuple[Dataset, Dataset]: + train_transform = transforms.Compose([ + transforms.RandomHorizontalFlip(dataset_config.hflip_prob), + transforms.Resize(dataset_config.train_size), + transforms.ToTensor(), + ]) + test_transform = transforms.Compose([ + transforms.Resize(dataset_config.test_size), + transforms.CenterCrop(dataset_config.test_crop_size), + transforms.ToTensor(), + ]) + if dataset_config.dataset in {'cifar10', 'cifar100', 'cifar'}: + if dataset_config.dataset in {'cifar10', 'cifar'}: + train_set = CIFAR10(dataset_config.dataset_dir, train=True, + transform=train_transform, download=True) + test_set = CIFAR10(dataset_config.dataset_dir, train=False, + transform=test_transform) + else: # CIFAR-100 + train_set = CIFAR100(dataset_config.dataset_dir, train=True, + transform=train_transform, download=True) + test_set = CIFAR100(dataset_config.dataset_dir, train=False, + transform=test_transform) + elif dataset_config.dataset in {'imagenet1k', 'imagenet'}: + train_set = ImageNet(dataset_config.dataset_dir, 'train', + transform=train_transform) + test_set = ImageNet(dataset_config.dataset_dir, 'val', + transform=test_transform) + else: + raise NotImplementedError(f"Unimplemented dataset: '{dataset_config.dataset}") + + return train_set, test_set + + def _init_models(self, dataset: str) -> Iterable[tuple[str, torch.nn.Module]]: + if dataset in {'cifar10', 'cifar100', 'cifar'}: + backbone = CIFARSimCLRResNet50(self.hid_dim, pretrain=False) + if dataset in {'cifar10', 'cifar'}: + classifier = torch.nn.Linear(self.hid_dim, 10) + else: + classifier = torch.nn.Linear(self.hid_dim, 100) + elif dataset in {'imagenet1k', 'imagenet'}: + backbone = ImageNetSimCLRResNet50(self.hid_dim, pretrain=False) + classifier = torch.nn.Linear(self.hid_dim, 1000) + else: + raise NotImplementedError(f"Unimplemented dataset: '{dataset}") + + yield 'backbone', backbone + yield 'classifier', classifier + + def _custom_init_fn(self, config: SimCLREvalConfig): + self.optims = {n: LARS(o) if config.optim_config.optim == 'lars' else o + for n, o in self.optims.items()} + if self.restore_iter == 0: + pretrained_checkpoint = torch.load(self.pretrained_checkpoint) + backbone_checkpoint = pretrained_checkpoint['model_state_dict'] + backbone_state_dict = {k: v for k, v in backbone_checkpoint.items() + if k in self.models['backbone'].state_dict()} + self.models['backbone'].load_state_dict(backbone_state_dict) + + def train(self, num_iters: int, loss_fn: Callable, logger: Loggers, device: torch.device): + backbone, classifier = self.models.values() + optim_b, optim_c = self.optims.values() + sched_b, sched_c = self.scheds.values() + loader_size = len(self.train_loader) + num_batches = num_iters * loader_size + for iter_ in range(self.restore_iter, num_iters): + if self.finetune: + backbone.train() + else: + backbone.eval() + classifier.train() + for batch, (images, targets) in enumerate(self.train_loader): + global_batch = iter_ * loader_size + batch + images, targets = images.to(device), targets.to(device) + classifier.zero_grad() + if self.finetune: + backbone.zero_grad() + embedding = backbone(images) + else: + with torch.no_grad(): + embedding = backbone(images) + logits = classifier(embedding) + train_loss = loss_fn(logits, targets) + train_loss.backward() + if self.finetune: + optim_b.step() + optim_c.step() + self.log(logger, self.BatchLogRecord( + batch, num_batches, global_batch, iter_, num_iters, + optim_c.param_groups[0]['lr'], train_loss.item() + )) + metrics = torch.Tensor(list(self.eval(loss_fn, device))).mean(0) + eval_loss = metrics[0].item() + eval_accuracy = metrics[1].item() + epoch_log = self.EpochLogRecord(iter_, num_iters, eval_loss, eval_accuracy) + self.log(logger, epoch_log) + self.save_checkpoint(epoch_log) + if sched_b is not None and self.finetune: + sched_b.step() + if sched_c is not None: + sched_c.step() + + def eval(self, loss_fn: Callable, device: torch.device): + backbone, classifier = self.models.values() + backbone.eval() + classifier.eval() + with torch.no_grad(): + for images, targets in self.test_loader: + images, targets = images.to(device), targets.to(device) + embedding = backbone(images) + logits = classifier(embedding) + loss = loss_fn(logits, targets) + prediction = logits.argmax(1) + accuracy = (prediction == targets).float().mean() + yield loss.item(), accuracy.item() + + +if __name__ == '__main__': + args = parse_args_and_config() + config = SimCLREvalConfig.from_args(args) + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + trainer = SimCLREvalTrainer( + seed=args.seed, + checkpoint_dir=args.checkpoint_dir, + device=device, + inf_mode=False, + num_iters=args.num_iters, + config=config, + hid_dim=args.hid_dim, + out_dim=args.out_dim, + pretrained_checkpoint=args.pretrained_checkpoint, + finetune=args.finetune, + ) + + loggers = trainer.init_logger(args.log_dir) + trainer.train(args.num_iters, torch.nn.CrossEntropyLoss(), loggers, device) diff --git a/simclr/models.py b/simclr/models.py index 9996188..8b270b9 100644 --- a/simclr/models.py +++ b/simclr/models.py @@ -7,17 +7,19 @@ from torchvision.models.resnet import BasicBlock # TODO Make a SimCLR base class class CIFARSimCLRResNet50(ResNet): - def __init__(self, hid_dim, out_dim): + def __init__(self, hid_dim, out_dim=128, pretrain=True): super(CIFARSimCLRResNet50, self).__init__( block=BasicBlock, layers=[3, 4, 6, 3], num_classes=hid_dim ) + self.pretrain = pretrain self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) - self.projector = nn.Sequential( - nn.Linear(hid_dim, hid_dim), - nn.ReLU(inplace=True), - nn.Linear(hid_dim, out_dim), - ) + if pretrain: + self.projector = nn.Sequential( + nn.Linear(hid_dim, hid_dim), + nn.ReLU(inplace=True), + nn.Linear(hid_dim, out_dim), + ) def backbone(self, x: Tensor) -> Tensor: x = self.conv1(x) @@ -37,22 +39,30 @@ class CIFARSimCLRResNet50(ResNet): def forward(self, x: Tensor) -> Tensor: h = self.backbone(x) - z = self.projector(h) - return z + if self.pretrain: + z = self.projector(h) + return z + else: + return h class ImageNetSimCLRResNet50(ResNet): - def __init__(self, hid_dim, out_dim): + def __init__(self, hid_dim, out_dim=128, pretrain=True): super(ImageNetSimCLRResNet50, self).__init__( block=BasicBlock, layers=[3, 4, 6, 3], num_classes=hid_dim ) - self.projector = nn.Sequential( - nn.Linear(hid_dim, hid_dim), - nn.ReLU(inplace=True), - nn.Linear(hid_dim, out_dim), - ) + self.pretrain = pretrain + if pretrain: + self.projector = nn.Sequential( + nn.Linear(hid_dim, hid_dim), + nn.ReLU(inplace=True), + nn.Linear(hid_dim, out_dim), + ) def forward(self, x: Tensor) -> Tensor: h = super(ImageNetSimCLRResNet50, self).forward(x) - z = self.projector(h) - return z + if self.pretrain: + z = self.projector(h) + return z + else: + return h |