Upload 19 files

base/__init__.py

@@ -0,0 +1,4 @@
+from .base_data_loader import *
+from .base_dataset import *
+from .base_model import *
+from .base_trainer import *

base/base_data_loader.py

@@ -0,0 +1,62 @@
+import numpy as np
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import default_collate
+from torch.utils.data.sampler import SubsetRandomSampler
+
+
+class BaseDataLoader(DataLoader):
+    """
+    Base class for all data loaders
+    """
+
+    def __init__(self, dataset, batch_size, shuffle, validation_split, num_workers, collate_fn=default_collate):
+        self.validation_split = validation_split
+        self.shuffle = shuffle
+
+        self.batch_idx = 0
+        self.n_samples = len(dataset)
+
+        self.sampler, self.valid_sampler = self._split_sampler(self.validation_split)
+
+        self.init_kwargs = {
+            'dataset': dataset,
+            'batch_size': batch_size,
+            'shuffle': self.shuffle,
+            'collate_fn': collate_fn,
+            'num_workers': num_workers
+        }
+        super().__init__(sampler=self.sampler, **self.init_kwargs)
+
+    def _split_sampler(self, split):
+        if split == 0.0:
+            return None, None
+
+        idx_full = np.arange(self.n_samples)
+
+        np.random.seed(0)
+        np.random.shuffle(idx_full)
+
+        if isinstance(split, int):
+            assert split > 0
+            assert split < self.n_samples, "validation set size is configured to be larger than entire dataset."
+            len_valid = split
+        else:
+            len_valid = int(self.n_samples * split)
+
+        valid_idx = idx_full[0:len_valid]
+        train_idx = np.delete(idx_full, np.arange(0, len_valid))
+
+        train_sampler = SubsetRandomSampler(train_idx)
+        valid_sampler = SubsetRandomSampler(valid_idx)
+
+        # turn off shuffle option which is mutually exclusive with sampler
+        self.shuffle = False
+        self.n_samples = len(train_idx)
+
+        return train_sampler, valid_sampler
+
+    def split_validation(self):
+        if self.valid_sampler is None:
+            return None
+        else:
+            return DataLoader(sampler=self.valid_sampler, **self.init_kwargs)

base/base_dataset.py

@@ -0,0 +1,35 @@
+import glob
+from pathlib import Path
+from typing import Any, Callable, Optional
+from torchvision.datasets import VisionDataset
+
+
+class BaseDataset(VisionDataset):
+    def __init__(
+        self,
+        root: str,
+        loader: Callable[[str], Any],
+        transforms: Optional[Callable] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        train: bool = True
+    ) -> None:
+        super().__init__(root, transforms, transform, target_transform)
+
+        self.root_path = Path(root)
+        self.loader = loader
+
+        mode = 'train' if train else 'test'
+        self.data = sorted(glob.glob(f'{mode}/images/*.jpg', root_dir=root))
+        self.masks = sorted(glob.glob(f'{mode}/masks/*.png', root_dir=root))
+
+    def __getitem__(self, index: int) -> Any:
+        img_path, mask_path = self.data[index], self.masks[index]
+        img_path, mask_path = self.root_path / img_path, self.root_path / mask_path
+
+        img, mask = self.loader(img_path), self.loader(mask_path)
+        img, mask = self.transforms(img, mask)
+        return img, mask.squeeze(dim=0).bool().float()
+
+    def __len__(self) -> int:
+        return len(self.data)

base/base_model.py

@@ -0,0 +1,25 @@
+from abc import abstractmethod
+import numpy as np
+import torch.nn as nn
+
+
+class BaseModel(nn.Module):
+    """
+    Base class for all models
+    """
+    @abstractmethod
+    def forward(self, *inputs):
+        """
+        Forward pass logic
+
+        :return: Model output
+        """
+        raise NotImplementedError
+
+    def __str__(self):
+        """
+        Model prints with number of trainable parameters
+        """
+        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
+        params = sum([np.prod(p.size()) for p in model_parameters])
+        return super().__str__() + '\nTrainable parameters: {}'.format(params)

base/base_trainer.py

@@ -0,0 +1,152 @@
+from abc import abstractmethod
+import torch
+from logger import TensorboardWriter
+from numpy import inf
+
+
+class BaseTrainer:
+    """
+    Base class for all trainers
+    """
+
+    def __init__(self, model, criterion, metric_ftns, optimizer, config):
+        self.config = config
+        self.logger = config.get_logger('trainer', config['trainer']['verbosity'])
+
+        self.model = model
+        self.criterion = criterion
+        self.metric_ftns = metric_ftns
+        self.optimizer = optimizer
+
+        cfg_trainer = config['trainer']
+        self.epochs = cfg_trainer['epochs']
+        self.save_period = cfg_trainer['save_period']
+        self.monitor = cfg_trainer.get('monitor', 'off')
+
+        # configuration to monitor model performance and save best
+        if self.monitor == 'off':
+            self.mnt_mode = 'off'
+            self.mnt_best = 0
+        else:
+            self.mnt_mode, self.mnt_metric = self.monitor.split()
+            assert self.mnt_mode in ['min', 'max']
+
+            self.mnt_best = inf if self.mnt_mode == 'min' else -inf
+            self.early_stop = cfg_trainer.get('early_stop', inf)
+            if self.early_stop <= 0:
+                self.early_stop = inf
+
+        self.start_epoch = 1
+
+        self.checkpoint_dir = config.save_dir
+
+        # setup visualization writer instance
+        self.writer = TensorboardWriter(config.log_dir, self.logger, cfg_trainer['tensorboard'])
+
+        if config.resume is not None:
+            self._resume_checkpoint(config.resume)
+
+    @abstractmethod
+    def _train_epoch(self, epoch):
+        """
+        Training logic for an epoch
+
+        :param epoch: Current epoch number
+        """
+        raise NotImplementedError
+
+    def train(self):
+        """
+        Full training logic
+        """
+        not_improved_count = 0
+        for epoch in range(self.start_epoch, self.epochs + 1):
+            result = self._train_epoch(epoch)
+
+            # save logged informations into log dict
+            log = {'epoch': epoch}
+            log.update(result)
+
+            # print logged informations to the screen
+            for key, value in log.items():
+                self.logger.info('    {:15s}: {}'.format(str(key), value))
+
+            # evaluate model performance according to configured metric, save best checkpoint as model_best
+            best = False
+            if self.mnt_mode != 'off':
+                try:
+                    # check whether model performance improved or not, according to specified metric(mnt_metric)
+                    improved = (self.mnt_mode == 'min' and log[self.mnt_metric] <= self.mnt_best) or \
+                               (self.mnt_mode == 'max' and log[self.mnt_metric] >= self.mnt_best)
+                except KeyError:
+                    self.logger.warning("Warning: Metric '{}' is not found. "
+                                        "Model performance monitoring is disabled.".format(self.mnt_metric))
+                    self.mnt_mode = 'off'
+                    improved = False
+
+                if improved:
+                    self.mnt_best = log[self.mnt_metric]
+                    not_improved_count = 0
+                    best = True
+                else:
+                    not_improved_count += 1
+
+                if not_improved_count > self.early_stop:
+                    self.logger.info("Validation performance didn\'t improve for {} epochs. "
+                                     "Training stops.".format(self.early_stop))
+                    break
+
+            if epoch % self.save_period == 0:
+                self._save_checkpoint(epoch, save_best=best)
+
+    def _save_checkpoint(self, epoch, save_best=False):
+        """
+        Saving checkpoints
+
+        :param epoch: current epoch number
+        :param log: logging information of the epoch
+        :param save_best: if True, rename the saved checkpoint to 'model_best.pth'
+        """
+        arch = type(self.model).__name__
+        state = {
+            'arch': arch,
+            'epoch': epoch,
+            'state_dict': self.model.state_dict(),
+            'optimizer': self.optimizer.state_dict(),
+            'monitor_best': self.mnt_best,
+            'config': self.config
+        }
+        filename = str(self.checkpoint_dir / 'checkpoint-epoch{}.pth'.format(epoch))
+        torch.save(state, filename)
+        self.logger.info("Saving checkpoint: {} ...".format(filename))
+        if save_best:
+            best_path = str(self.checkpoint_dir / 'model_best.pth')
+            torch.save(state, best_path)
+            self.logger.info("Saving current best: model_best.pth ...")
+
+    def _resume_checkpoint(self, resume_path):
+        """
+        Resume from saved checkpoints
+
+        :param resume_path: Checkpoint path to be resumed
+        """
+        resume_path = str(resume_path)
+        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
+        checkpoint = torch.load(resume_path)
+        self.start_epoch = checkpoint['epoch'] + 1
+        self.mnt_best = checkpoint['monitor_best']
+
+        # load architecture params from checkpoint.
+        if checkpoint['config']['arch'] != self.config['arch']:
+            self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
+                                "checkpoint. This may yield an exception while state_dict is being loaded.")
+        self.model.load_state_dict(checkpoint['state_dict'])
+
+        # load optimizer state from checkpoint only when optimizer type is not changed.
+        if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
+            self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
+                                "Optimizer parameters not being resumed.")
+        else:
+            self.optimizer.load_state_dict(checkpoint['optimizer'])
+
+        self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch))

data_loader/data_loaders.py

@@ -0,0 +1,83 @@
+from base import BaseDataLoader
+from dataset.datasets import PatchedDataset
+from torchvision import transforms
+from torch.utils.data.sampler import SequentialSampler
+
+
+class PatchedDataLoader(BaseDataLoader):
+    def __init__(
+        self,
+        data_dir,
+        patch_size,
+        batch_size,
+        patch_stride=None,
+        preds=None,
+        target_dist=0.0,
+        shuffle=True,
+        validation_split=0.0,
+        num_workers=1,
+        training=True
+    ):
+        trsfm = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.3551, 0.4698, 0.2261),
+                                 (0.1966, 0.1988, 0.1761))
+        ])
+        target_trsfm = transforms.Compose([
+            transforms.ToTensor(),
+        ])
+        rand_trsfm = transforms.RandomApply([
+            transforms.RandomVerticalFlip(),
+            transforms.RandomHorizontalFlip()
+        ])
+        self.data_dir = data_dir
+        self.dataset = PatchedDataset(
+            self.data_dir,
+            patch_size,
+            patch_stride=patch_stride,
+            preds=preds,
+            target_dist=target_dist,
+            transform=trsfm,
+            target_transform=target_trsfm,
+            rand_transform=rand_trsfm if training and shuffle else None,
+            train=training,
+            late_init=True
+        )
+        super().__init__(self.dataset, batch_size, shuffle, validation_split, num_workers)
+
+    def _split_sampler(self, split):
+        train_sampler, valid_sampler = super()._split_sampler(split)
+
+        if valid_sampler is not None:
+            self.dataset.make_dataset(valid_indices=valid_sampler.indices)
+        else:
+            self.dataset.make_dataset()
+
+        train_idx, valid_idx = [], []
+        for patch in self.dataset.patches:
+            if valid_sampler is not None and patch.idx in valid_sampler.indices:
+                valid_idx.append(self.dataset.patches.index(patch))
+            else:
+                train_idx.append(self.dataset.patches.index(patch))
+
+        if valid_sampler is not None:
+            train_sampler.indices, valid_sampler.indices = train_idx, valid_idx
+        else:
+            train_sampler = SequentialSampler(train_idx)
+
+        # turn off shuffle option which is mutually exclusive with sampler
+        self.shuffle = False
+        self.n_samples = len(train_idx)
+
+        return train_sampler, valid_sampler
+
+    def update_dataset(self, preds):
+        self.dataset.preds = preds
+        self.dataset.patches.clear()
+        self.n_samples = len(self.dataset)
+
+        train_sampler, valid_sampler = self._split_sampler(
+            self.validation_split)
+        if valid_sampler is not None:
+            self.valid_sampler.indices = valid_sampler.indices
+        self.sampler.indices = train_sampler.indices

dataset/datasets.py

@@ -0,0 +1,60 @@
+from typing import Any, Callable, Optional
+import torch.nn.functional as F
+from base import BaseDataset
+from utils.util import TransformMultiple, pil_loader
+from dataset.patches import Patches
+
+
+class PatchedDataset(BaseDataset):
+    def __init__(
+        self,
+        root: str,
+        patch_size: int,
+        patch_stride: int = None,
+        preds: list = None,
+        target_dist: float = 0.0,
+        transforms: Optional[Callable] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        rand_transform: Optional[Callable] = None,
+        train: bool = True,
+        late_init: bool = False
+    ) -> None:
+        super().__init__(root, pil_loader, transforms, transform, target_transform, train)
+        self.patches = Patches(patch_size, patch_stride)
+        self.preds = preds
+        self.target_dist = target_dist * patch_size ** 2
+        self.rand_transform = TransformMultiple(rand_transform)
+        if not late_init:
+            self.make_dataset()
+
+    def make_dataset(self, valid_indices=[]):
+        for idx in range(super().__len__()):
+            _, mask = super().__getitem__(idx)
+            if self.preds is not None:
+                mask = self._union_mask(mask, self.preds[idx])
+
+            if idx not in valid_indices:
+                self.patches.create(idx, mask, cond_fn=self._dist_fn
+                                    if self.target_dist != 0.0 else None)
+            else:
+                self.patches.create(idx, mask, no_overlap=True)
+
+    def __getitem__(self, index: int) -> Any:
+        patch = self.patches[index]
+        img, mask = super().__getitem__(patch.idx)
+
+        img_patch = self.patches.get_patch(img, patch)
+        mask_patch = self.patches.get_patch(mask, patch)
+        img_patch, mask_patch = self.rand_transform(
+            (img_patch, mask_patch.unsqueeze(dim=0)))
+        return img_patch, mask_patch.squeeze(dim=0)
+
+    def _union_mask(self, mask, pred):
+        pred = F.pad(
+            pred, (0, mask.shape[1] - pred.shape[1], 0, mask.shape[0] - pred.shape[0]))
+        return (mask + pred) - (mask * pred)
+
+    def _dist_fn(self, mask, patch):
+        data = self.patches.get_patch(mask, patch)
+        return data.count_nonzero() > self.target_dist

dataset/patches.py

@@ -0,0 +1,49 @@
+from collections import UserList
+from torchvision.utils import make_grid
+
+
+class Patch():
+    def __init__(self, idx, x, y) -> None:
+        self.idx = idx
+        self.x, self.y = x, y
+        self.data = None
+
+    def __eq__(self, __o: object) -> bool:
+        return self.idx == __o.idx and self.x == __o.x and self.y == __o.y
+
+
+class Patches(UserList):
+    def __init__(self, size, stride=None):
+        super().__init__()
+        self.size = size
+        self.stride = stride if stride is not None else size
+
+    def create(self, index, data, cond_fn=None, no_overlap=False):
+        stride = self.size if no_overlap else self.stride
+        for x in range(0, data.size(-2) - self.size + 1, stride):
+            for y in range(0, data.size(-1) - self.size + 1, stride):
+                patch = Patch(index, x, y)
+                if cond_fn is None or cond_fn(data, patch):
+                    self.append(patch)
+
+    def get_patch(self, data, patch: Patch):
+        assert data.ndim in {2, 3}, 'only 2-D and 3-D Tensors are supported.'
+        _data = data.unsqueeze(dim=0) if data.ndim == 2 else data
+        data_patch = _data[:, patch.x:patch.x + self.size,
+                           patch.y:patch.y + self.size]
+        return data_patch.squeeze(dim=0) if data.ndim == 2 else data_patch
+
+    def store_data(self, indices, data):
+        for idx in range(len(indices)):
+            self[indices[idx]].data = [data[i][idx] for i in range(len(data))]
+
+    def retrieve_data(self, indices):
+        return [[self[idx].data[i] for idx in indices] for i in range(len(self[indices[0]].data))]
+
+    def combine(self, index: int, data_idx: int):
+        indices = [self.index(patch) for patch in self if patch.idx == index]
+        indices.sort(key=lambda idx: (self[idx].x, self[idx].y))
+
+        data = self.retrieve_data(indices)
+        nrow = sum([self[idx].x == 0 for idx in indices])
+        return make_grid(data[data_idx], nrow, padding=0)

logger/__init__.py

@@ -0,0 +1,2 @@
+from .logger import *
+from .visualization import *

logger/logger.py

@@ -0,0 +1,23 @@
+import logging
+import logging.config
+from pathlib import Path
+
+from utils import read_json
+
+
+def setup_logging(save_dir, log_config='logger/logger_config.json', default_level=logging.INFO):
+    """
+    Setup logging configuration
+    """
+    log_config = Path(log_config)
+    if log_config.is_file():
+        config = read_json(log_config)
+        # modify logging paths based on run config
+        for _, handler in config['handlers'].items():
+            if 'filename' in handler:
+                handler['filename'] = str(save_dir / handler['filename'])
+
+        logging.config.dictConfig(config)
+    else:
+        print("Warning: logging configuration file is not found in {}.".format(log_config))
+        logging.basicConfig(level=default_level)

logger/logger_config.json

@@ -0,0 +1,36 @@
+{
+    "version": 1,
+    "disable_existing_loggers": false,
+    "formatters": {
+        "simple": {
+            "format": "%(message)s"
+        },
+        "datetime": {
+            "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        }
+    },
+    "handlers": {
+        "console": {
+            "class": "logging.StreamHandler",
+            "level": "DEBUG",
+            "formatter": "simple",
+            "stream": "ext://sys.stdout"
+        },
+        "info_file_handler": {
+            "class": "logging.handlers.RotatingFileHandler",
+            "level": "INFO",
+            "formatter": "datetime",
+            "filename": "info.log",
+            "maxBytes": 10485760,
+            "backupCount": 20,
+            "encoding": "utf8"
+        }
+    },
+    "root": {
+        "level": "INFO",
+        "handlers": [
+            "console",
+            "info_file_handler"
+        ]
+    }
+}

logger/visualization.py

@@ -0,0 +1,78 @@
+import importlib
+from datetime import datetime
+
+
+class TensorboardWriter():
+    def __init__(self, log_dir, logger, enabled):
+        self.writer = None
+        self.selected_module = ""
+
+        if enabled:
+            log_dir = str(log_dir)
+
+            # Retrieve vizualization writer.
+            succeeded = False
+            for module in ["torch.utils.tensorboard", "tensorboardX"]:
+                try:
+                    self.writer = importlib.import_module(module).SummaryWriter(log_dir)
+                    succeeded = True
+                    break
+                except ImportError:
+                    succeeded = False
+                self.selected_module = module
+
+            if not succeeded:
+                message = "Warning: visualization (Tensorboard) is configured to use, but currently not installed on " \
+                    "this machine. Please install TensorboardX with 'pip install tensorboardx', upgrade PyTorch to " \
+                    "version >= 1.1 to use 'torch.utils.tensorboard' or turn off the option in the 'config.json' file."
+                logger.warning(message)
+
+        self.step = 0
+        self.mode = ''
+        self.step_tracker = {}
+
+        self.tb_writer_ftns = {
+            'add_scalar', 'add_scalars', 'add_image', 'add_images', 'add_audio',
+            'add_text', 'add_histogram', 'add_pr_curve', 'add_embedding'
+        }
+        self.tag_mode_exceptions = {'add_histogram', 'add_embedding'}
+        self.timer = datetime.now()
+
+    def set_step(self, step, mode='train'):
+        self.mode = mode
+        self.step = step
+        if step == 0:
+            self.timer = datetime.now()
+        else:
+            duration = datetime.now() - self.timer
+            self.add_scalar('steps_per_sec', 1 / duration.total_seconds())
+            self.timer = datetime.now()
+
+    def next(self, mode='train'):
+        step = self.step_tracker[mode] = self.step_tracker.get(mode, 0) + 1
+        self.set_step(step, mode=mode)
+
+    def __getattr__(self, name):
+        """
+        If visualization is configured to use:
+            return add_data() methods of tensorboard with additional information (step, tag) added.
+        Otherwise:
+            return a blank function handle that does nothing
+        """
+        if name in self.tb_writer_ftns:
+            add_data = getattr(self.writer, name, None)
+
+            def wrapper(tag, data, *args, **kwargs):
+                if add_data is not None:
+                    # add mode(train/valid) tag
+                    if name not in self.tag_mode_exceptions:
+                        tag = '{}/{}'.format(tag, self.mode)
+                    add_data(tag, data, self.step, *args, **kwargs)
+            return wrapper
+        else:
+            # default action for returning methods defined in this class, set_step() for instance.
+            try:
+                attr = object.__getattr__(name)
+            except AttributeError:
+                raise AttributeError("type object '{}' has no attribute '{}'".format(self.selected_module, name))
+            return attr

model/loss.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn.functional as F
+
+
+def focal_loss(inputs, targets, alpha=0.5, gamma=2, reduction='mean'):
+    logpt = F.cross_entropy(inputs, targets.long(), reduction='none')
+    pt = torch.exp(-logpt)
+    focal_loss = (1 - pt) ** gamma * logpt
+    alpha_weight = alpha * targets + (1 - alpha) * (1 - targets)
+    focal_loss = alpha_weight * focal_loss
+
+    if reduction == 'mean':
+        return torch.mean(focal_loss)
+    elif reduction == 'sum':
+        return torch.sum(focal_loss)
+    else:
+        return focal_loss
+
+
+def dice_loss(inputs, targets, epsilon=1e-7):
+    targets_one_hot = torch.nn.functional.one_hot(targets.long(), num_classes=inputs.shape[1])
+    targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float()
+    inputs = F.softmax(inputs, dim=1)
+    targets_one_hot = targets_one_hot.type(inputs.type())
+    numerator = 2 * (inputs * targets_one_hot).sum(dim=(2,3))
+    denominator = inputs.sum(dim=(2,3)) + targets_one_hot.sum(dim=(2,3))
+    dice_coefficient = numerator / (denominator + epsilon)
+    return 1 - dice_coefficient.mean()

model/metric.py

@@ -0,0 +1,24 @@
+import torch
+from sklearn.metrics import f1_score as f1
+from sklearn.metrics import precision_score, recall_score
+
+
+def precision(output, target):
+    with torch.no_grad():
+        pred = torch.argmax(output, dim=1)
+        assert pred.shape[0] == len(target)
+    return precision_score(target.view(-1).cpu(), pred.view(-1).cpu())
+
+
+def recall(output, target):
+    with torch.no_grad():
+        pred = torch.argmax(output, dim=1)
+        assert pred.shape[0] == len(target)
+    return recall_score(target.view(-1).cpu(), pred.view(-1).cpu())
+
+
+def f1_score(output, target):
+    with torch.no_grad():
+        pred = torch.argmax(output, dim=1)
+        assert pred.shape[0] == len(target)
+    return f1(target.view(-1).cpu(), pred.view(-1).cpu())

model/model.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from base import BaseModel
+
+
+class UNet(BaseModel):
+    def __init__(self, n_channels, n_classes, bilinear=False):
+        super(UNet, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+
+        self.inc = DoubleConv(n_channels, 64)
+        self.down1 = Down(64, 128)
+        self.down2 = Down(128, 256)
+        self.down3 = Down(256, 512)
+        factor = 2 if bilinear else 1
+        self.down4 = Down(512, 1024 // factor)
+        self.up1 = Up(1024, 512 // factor, bilinear)
+        self.up2 = Up(512, 256 // factor, bilinear)
+        self.up3 = Up(256, 128 // factor, bilinear)
+        self.up4 = Up(128, 64, bilinear)
+        self.outc = OutConv(64, n_classes)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        logits = self.outc(x)
+        return F.log_softmax(logits, dim=1)
+
+
+class DoubleConv(BaseModel):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class Down(BaseModel):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(kernel_size=8, stride=2, padding=3),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(BaseModel):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class OutConv(BaseModel):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)

trainer/__init__.py

@@ -0,0 +1 @@
+from .trainer import *

trainer/trainer.py

@@ -0,0 +1,135 @@
+import numpy as np
+import torch
+from base import BaseTrainer
+from torchvision.utils import make_grid
+from utils import MetricTracker
+
+
+class Trainer(BaseTrainer):
+    """
+    Trainer class
+    """
+
+    def __init__(self, model, criterion, metric_ftns, optimizer, config, device,
+                 data_loader, valid_data_loader=None, lr_scheduler=None, len_epoch=None):
+        super().__init__(model, criterion, metric_ftns, optimizer, config)
+        self.config = config
+        self.device = device
+        self.data_loader = data_loader
+        self.len_epoch = len(self.data_loader)
+        self.valid_data_loader = valid_data_loader
+        self.do_validation = self.valid_data_loader is not None
+        self.lr_scheduler = lr_scheduler
+        self.log_step = int(np.sqrt(data_loader.batch_size))
+        self.adaptive_step = config['trainer']['adaptive_step']
+
+        self.train_metrics = MetricTracker(
+            'loss', *[m.__name__ for m in self.metric_ftns], writer=self.writer)
+        self.valid_metrics = MetricTracker(
+            'loss', *[m.__name__ for m in self.metric_ftns], writer=self.writer)
+
+    def _train_epoch(self, epoch):
+        """
+        Training logic for an epoch
+
+        :param epoch: Integer, current training epoch.
+        :return: A log that contains average loss and metric in this epoch.
+        """
+        if epoch > self.adaptive_step and epoch % self.adaptive_step == 1:
+            dataset = self.data_loader.inference.dataset
+            self.model.eval()
+            with torch.no_grad():
+                for batch_idx, (data, target) in enumerate(self.data_loader.inference):
+                    data, target = data.to(self.device), target.to(self.device)
+                    output = self.model(data)
+
+                    batch_size = self.data_loader.inference.batch_size
+                    patch_idx = torch.arange(
+                        batch_size * batch_idx, batch_size * batch_idx + data.shape[0])
+                    pred = torch.argmax(output, dim=1)
+                    dataset.patches.store_data(patch_idx, [pred.unsqueeze(1)])
+
+            preds = [dataset.patches.combine(idx, data_idx=0)[0].cpu()
+                    for idx in range(len(dataset.data))]
+
+            self.data_loader.update_dataset(preds)
+            self.len_epoch = len(self.data_loader)
+
+        self.model.train()
+        self.train_metrics.reset()
+        for batch_idx, (data, target) in enumerate(self.data_loader):
+            data, target = data.to(self.device), target.to(self.device)
+
+            self.optimizer.zero_grad()
+            output = self.model(data)
+            loss = self.criterion(output, target)
+            loss.backward()
+            self.optimizer.step()
+
+            self.train_metrics.update('loss', loss.item())
+            for met in self.metric_ftns:
+                self.train_metrics.update(met.__name__, met(output, target))
+
+            if batch_idx % self.log_step == 0:
+                self.logger.debug('Train Epoch: {} {} Loss: {:.6f}'.format(
+                    epoch,
+                    self._progress(batch_idx),
+                    loss.item()))
+                self.writer.add_image('input', make_grid(
+                    data.cpu(), nrow=8, normalize=True))
+
+            if batch_idx == self.len_epoch:
+                break
+        
+        self.writer.next()
+        self.train_metrics.add_scalers()
+        log = self.train_metrics.result()
+
+        if self.do_validation:
+            val_log = self._valid_epoch(epoch)
+            log.update(**{'val_' + k: v for k, v in val_log.items()})
+
+        if self.lr_scheduler is not None:
+            self.lr_scheduler.step()
+        return log
+
+    def _valid_epoch(self, epoch):
+        """
+        Validate after training an epoch
+
+        :param epoch: Integer, current training epoch.
+        :return: A log that contains information about validation
+        """
+        self.model.eval()
+        self.valid_metrics.reset()
+        with torch.no_grad():
+            for batch_idx, (data, target) in enumerate(self.valid_data_loader):
+                data, target = data.to(self.device), target.to(self.device)
+
+                output = self.model(data)
+                loss = self.criterion(output, target)
+
+                self.valid_metrics.update('loss', loss.item())
+                for met in self.metric_ftns:
+                    self.valid_metrics.update(
+                        met.__name__, met(output, target))
+                self.writer.add_image('input', make_grid(
+                    data.cpu(), nrow=8, normalize=True))
+
+        # add histogram of model parameters to the tensorboard
+        for name, p in self.model.named_parameters():
+            self.writer.add_histogram(name, p, bins='auto')
+        
+        self.writer.next('valid')
+        self.valid_metrics.add_scalers()
+        return self.valid_metrics.result()
+
+    def _progress(self, batch_idx):
+        base = '[{}/{} ({:.0f}%)]'
+        if hasattr(self.data_loader, 'n_samples'):
+            current = batch_idx * self.data_loader.batch_size
+            total = self.data_loader.n_samples
+        else:
+            current = batch_idx
+            total = self.len_epoch
+        return base.format(current, total, 100.0 * current / total)

utils/__init__.py

@@ -0,0 +1 @@
+from .util import *

utils/util.py

@@ -0,0 +1,91 @@
+import json
+from collections import OrderedDict
+from itertools import repeat
+from pathlib import Path
+from typing import Any, Callable, List, Optional, Tuple
+import pandas as pd
+import PIL.Image as Image
+import torch
+
+
+def read_json(fname):
+    fname = Path(fname)
+    with fname.open('rt') as handle:
+        return json.load(handle, object_hook=OrderedDict)
+
+
+def write_json(content, fname):
+    fname = Path(fname)
+    with fname.open('wt') as handle:
+        json.dump(content, handle, indent=4, sort_keys=False)
+
+
+def pil_loader(fname: str) -> Image.Image:
+    return Image.open(fname)
+
+
+def prepare_device(n_gpu_use):
+    """
+    setup GPU device if available. get gpu device indices which are used for DataParallel
+    """
+    n_gpu = torch.cuda.device_count()
+    if n_gpu_use > 0 and n_gpu == 0:
+        print("Warning: There\'s no GPU available on this machine,"
+              "training will be performed on CPU.")
+        n_gpu_use = 0
+    if n_gpu_use > n_gpu:
+        print(f"Warning: The number of GPU\'s configured to use is {n_gpu_use}, but only {n_gpu} are "
+              "available on this machine.")
+        n_gpu_use = n_gpu
+    device = torch.device('cuda:0' if n_gpu_use > 0 else 'cpu')
+    list_ids = list(range(n_gpu_use))
+    return device, list_ids
+
+
+class TransformMultiple:
+    def __init__(self, transform: Optional[Callable] = None) -> None:
+        self.transform = transform
+
+    def __call__(self, data: Any) -> Tuple:
+        if self.transform is not None:
+            cat_data = torch.cat(data)
+            cat_data = self.transform(cat_data)
+            return torch.split(cat_data, [x.size()[0] for x in data])
+        return data
+
+    def _format_transform_repr(self, transform: Callable, head: str) -> List[str]:
+        lines = transform.__repr__().splitlines()
+        return [f"{head}{lines[0]}"] + ["{}{}".format(" " * len(head), line) for line in lines[1:]]
+
+    def __repr__(self) -> str:
+        body = [self.__class__.__name__]
+        if self.transform is not None:
+            body += self._format_transform_repr(self.transform, "Transform: ")
+        return "\n".join(body)
+
+
+class MetricTracker:
+    def __init__(self, *keys, writer=None):
+        self.writer = writer
+        self._data = pd.DataFrame(index=keys, columns=['total', 'counts', 'average'])
+        self.reset()
+
+    def reset(self):
+        for col in self._data.columns:
+            self._data[col].values[:] = 0
+
+    def update(self, key, value, n=1):
+        self._data.total[key] += value * n
+        self._data.counts[key] += n
+        self._data.average[key] = self._data.total[key] / self._data.counts[key]
+
+    def add_scalers(self):
+        if self.writer is not None:
+            for key in self._data.index:
+                self.writer.add_scalar(key, self._data.average[key])
+
+    def avg(self, key):
+        return self._data.average[key]
+
+    def result(self):
+        return dict(self._data.average)