add code for few-shot baseline

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]
+
+
+
+

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}")

model/__init__.py

@@ -1,2 +1,3 @@
 from model.registry import MODEL
 import model.residual_generator
+import model.fewshot

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)