From 8102651a282f1c6f5d433764b4956a855fd8886c Mon Sep 17 00:00:00 2001 From: Ray Wong Date: Mon, 10 Aug 2020 08:51:26 +0800 Subject: [PATCH] add code for few-shot baseline --- configs/few-shot/crossdomain.yml | 62 ++++++++++++++++++ engine/crossdomain.py | 97 ++++++++++++++++++++++++++++ model/__init__.py | 1 + model/fewshot.py | 105 +++++++++++++++++++++++++++++++ 4 files changed, 265 insertions(+) create mode 100644 configs/few-shot/crossdomain.yml create mode 100644 engine/crossdomain.py create mode 100644 model/fewshot.py diff --git a/configs/few-shot/crossdomain.yml b/configs/few-shot/crossdomain.yml new file mode 100644 index 0000000..8b4cb56 --- /dev/null +++ b/configs/few-shot/crossdomain.yml @@ -0,0 +1,62 @@ +name: cross-domain +engine: crossdomain +result_dir: ./result + +distributed: + model: + # broadcast_buffers: False + +misc: + random_seed: 1004 + +checkpoints: + interval: 2000 + +log: + logger: + level: 20 # DEBUG(10) INFO(20) + +model: + _type: resnet10 + +baseline: + plusplus: False + optimizers: + _type: Adam + data: + dataloader: + batch_size: 1024 + shuffle: True + num_workers: 16 + pin_memory: True + drop_last: True + dataset: + train: + path: /data/few-shot/mini_imagenet_full_size/train + pipeline: + - RandomResizedCrop: + size: [256, 256] + - ColorJitter: + brightness: 0.4 + contrast: 0.4 + saturation: 0.4 + - RandomHorizontalFlip + - ToTensor + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + val: + path: /data/few-shot/mini_imagenet_full_size/val + pipeline: + - Resize: + size: [286, 286] + - RandomCrop: + size: [256, 256] + - ToTensor + - Normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + + + + diff --git a/engine/crossdomain.py b/engine/crossdomain.py new file mode 100644 index 0000000..be3db60 --- /dev/null +++ b/engine/crossdomain.py @@ -0,0 +1,97 @@ +import torch +import torch.nn as nn +from torchvision.datasets import ImageFolder + +import ignite.distributed as idist +from ignite.contrib.metrics.gpu_info import GpuInfo +from ignite.contrib.handlers.tensorboard_logger import TensorboardLogger, global_step_from_engine, OutputHandler, \ + WeightsScalarHandler, GradsHistHandler, WeightsHistHandler, GradsScalarHandler +from ignite.engine import create_supervised_evaluator, create_supervised_trainer, Events +from ignite.metrics import Accuracy, Loss, RunningAverage +from ignite.contrib.engines.common import save_best_model_by_val_score +from ignite.contrib.handlers import ProgressBar + +from util.build import build_model, build_optimizer +from util.handler import setup_common_handlers +from data.transform import transform_pipeline + + +def baseline_trainer(config, logger, val_loader): + model = build_model(config.model, config.distributed.model) + optimizer = build_optimizer(model.parameters(), config.baseline.optimizers) + loss_fn = nn.CrossEntropyLoss() + trainer = create_supervised_trainer(model, optimizer, loss_fn, idist.device(), non_blocking=True) + trainer.logger = logger + RunningAverage(output_transform=lambda x: x).attach(trainer, "loss") + ProgressBar(ncols=0).attach(trainer) + + val_metrics = { + "accuracy": Accuracy(), + "nll": Loss(loss_fn) + } + evaluator = create_supervised_evaluator(model, val_metrics, idist.device()) + ProgressBar(ncols=0).attach(evaluator) + + @trainer.on(Events.EPOCH_COMPLETED) + def log_training_loss(engine): + logger.info(f"Epoch[{engine.state.epoch}] Loss: {engine.state.output:.2f}") + evaluator.run(val_loader) + metrics = evaluator.state.metrics + logger.info("Training Results - Avg accuracy: {:.2f} Avg loss: {:.2f}" + .format(trainer.state.epoch, metrics["accuracy"], metrics["nll"])) + + if idist.get_rank() == 0: + GpuInfo().attach(trainer, name='gpu') + + tb_logger = TensorboardLogger(log_dir=config.output_dir) + tb_logger.attach( + evaluator, + log_handler=OutputHandler( + tag="val", + metric_names='all', + global_step_transform=global_step_from_engine(trainer), + ), + event_name=Events.EPOCH_COMPLETED + ) + tb_logger.attach(trainer, log_handler=WeightsScalarHandler(model), + event_name=Events.EPOCH_COMPLETED(every=10)) + + tb_logger.attach(trainer, log_handler=WeightsHistHandler(model), event_name=Events.EPOCH_COMPLETED(every=25)) + + tb_logger.attach(trainer, log_handler=GradsScalarHandler(model), + event_name=Events.EPOCH_COMPLETED(every=10)) + + tb_logger.attach(trainer, log_handler=GradsHistHandler(model), event_name=Events.EPOCH_COMPLETED(every=25)) + + @trainer.on(Events.COMPLETED) + def _(): + tb_logger.close() + + to_save = dict(model=model, optimizer=optimizer, trainer=trainer) + setup_common_handlers(trainer, config.output_dir, print_interval_event=Events.EPOCH_COMPLETED, to_save=to_save, + save_interval_event=Events.EPOCH_COMPLETED(every=25), n_saved=5, + metrics_to_print=["loss"]) + save_best_model_by_val_score(config.output_dir, evaluator, model, "accuracy", 1, trainer) + return trainer + + +def run(task, config, logger): + assert torch.backends.cudnn.enabled + torch.backends.cudnn.benchmark = True + logger.info(f"start task {task}") + if task == "baseline": + train_dataset = ImageFolder(config.baseline.data.dataset.train.path, + transform=transform_pipeline(config.baseline.data.dataset.train.pipeline)) + val_dataset = ImageFolder(config.baseline.data.dataset.val.path, + transform=transform_pipeline(config.baseline.data.dataset.val.pipeline)) + logger.info(f"train with dataset:\n{train_dataset}") + train_data_loader = idist.auto_dataloader(train_dataset, **config.baseline.data.dataloader) + val_data_loader = idist.auto_dataloader(val_dataset, **config.baseline.data.dataloader) + trainer = baseline_trainer(config, logger, val_data_loader) + try: + trainer.run(train_data_loader, max_epochs=400) + except Exception: + import traceback + print(traceback.format_exc()) + else: + return NotImplemented(f"invalid task: {task}") diff --git a/model/__init__.py b/model/__init__.py index 28ea45c..328df68 100644 --- a/model/__init__.py +++ b/model/__init__.py @@ -1,2 +1,3 @@ from model.registry import MODEL import model.residual_generator +import model.fewshot diff --git a/model/fewshot.py b/model/fewshot.py new file mode 100644 index 0000000..851f6d9 --- /dev/null +++ b/model/fewshot.py @@ -0,0 +1,105 @@ +import math + +import torch.nn as nn + +from .registry import MODEL + + +# --- gaussian initialize --- +def init_layer(l): + # Initialization using fan-in + if isinstance(l, nn.Conv2d): + n = l.kernel_size[0] * l.kernel_size[1] * l.out_channels + l.weight.data.normal_(0, math.sqrt(2.0 / float(n))) + elif isinstance(l, nn.BatchNorm2d): + l.weight.data.fill_(1) + l.bias.data.fill_(0) + elif isinstance(l, nn.Linear): + l.bias.data.fill_(0) + + +class Flatten(nn.Module): + def __init__(self): + super(Flatten, self).__init__() + + def forward(self, x): + return x.view(x.size(0), -1) + + +class SimpleBlock(nn.Module): + def __init__(self, in_channels, out_channels, half_res, leakyrelu=False): + super(SimpleBlock, self).__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2 if half_res else 1, padding=1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace=True) if not leakyrelu else nn.LeakyReLU(0.2, inplace=True), + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False), + nn.BatchNorm2d(out_channels), + ) + self.relu = nn.ReLU(inplace=True) if not leakyrelu else nn.LeakyReLU(0.2, inplace=True) + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, 2 if half_res else 1, bias=False), + nn.BatchNorm2d(out_channels) + ) + else: + self.shortcut = nn.Identity() + + def forward(self, x): + o = self.block(x) + return self.relu(o + self.shortcut(x)) + + +class ResNet(nn.Module): + def __init__(self, block, layers, dims, num_classes=None, classifier_type="linear", flatten=True, leakyrelu=False): + super().__init__() + assert len(layers) == 4, 'Can have only four stages' + self.inplanes = 64 + + self.start = nn.Sequential( + nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False), + nn.BatchNorm2d(self.inplanes), + nn.ReLU(inplace=True) if not leakyrelu else nn.LeakyReLU(0.2, inplace=True), + nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + ) + + trunk = [] + in_channels = self.inplanes + for i in range(4): + for j in range(layers[i]): + half_res = i >= 1 and j == 0 + trunk.append(block(in_channels, dims[i], half_res, leakyrelu)) + in_channels = dims[i] + if flatten: + trunk.append(nn.AvgPool2d(7)) + trunk.append(Flatten()) + if num_classes is not None: + if classifier_type == "linear": + trunk.append(nn.Linear(in_channels, num_classes)) + elif classifier_type == "distlinear": + pass + else: + raise ValueError(f"invalid classifier_type:{classifier_type}") + self.trunk = nn.Sequential(*trunk) + self.apply(init_layer) + + def forward(self, x): + return self.trunk(self.start(x)) + + +@MODEL.register_module() +def resnet10(num_classes=None, classifier_type="linear", flatten=True, leakyrelu=False): + return ResNet(SimpleBlock, [1, 1, 1, 1], [64, 128, 256, 512], num_classes, classifier_type, flatten, leakyrelu) + + +@MODEL.register_module() +def resnet18(num_classes=None, classifier_type="linear", flatten=True, leakyrelu=False): + return ResNet(SimpleBlock, [2, 2, 2, 2], [64, 128, 256, 512], num_classes, classifier_type, flatten, leakyrelu) + + +@MODEL.register_module() +def resnet34(num_classes=None, classifier_type="linear", flatten=True, leakyrelu=False): + return ResNet(SimpleBlock, [3, 4, 6, 3], [64, 128, 256, 512], num_classes, classifier_type, flatten, leakyrelu)