[cnn_architecture]: Add starting blocks to experiment with

Rizhiy · Rizhiy · commit ebe6d752f204 · 2019-04-22T02:08:08.000+03:00
Currently, single run only
diff --git a/.gitignore b/.gitignore
@@ -8,3 +8,5 @@
 data
 *.pkl
 *.onnx
+*.model
+*.log
diff --git a/cnn_architecture/__init__.py b/cnn_architecture/__init__.py
diff --git a/cnn_architecture/blocks.py b/cnn_architecture/blocks.py
@@ -0,0 +1,117 @@
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+
+# Complexity of Conv2d, without feature_map_size: kernel_size^2*in_channels*out_channels
+
+# All blocks should have equal or lower computational complexity than simple block
+# Count multiplications only as an approximation
+# Total_complexity: 9*in*hidden + 3*9*hidden^2
+class SimpleBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self._num_hidden = in_channels * 2
+        self._conv1 = nn.Conv2d(in_channels, self._num_hidden, 3, stride=2, padding=1)
+        self._conv2 = nn.Conv2d(self._num_hidden, self._num_hidden, 3, padding=1)
+        self._conv3 = nn.Conv2d(self._num_hidden, self._num_hidden, 3, padding=1)
+
+    def forward(self, batch: Tensor) -> Tensor:
+        batch = F.relu(self._conv1(batch))
+        batch = F.relu(self._conv2(batch))
+        batch = F.relu(self._conv3(batch))
+        return batch
+
+    @property
+    def num_hidden(self):
+        return self._num_hidden
+
+
+# Total complexity: 9*in*hidden + 3*9*(hidden*hidden/2+hidden/2*hidden) + 3*relu
+class AddBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self._num_hidden = in_channels * 2
+        self._stride = nn.Conv2d(in_channels, self._num_hidden, 3, stride=2, padding=1)
+
+        self._conv1 = nn.Conv2d(self._num_hidden, self._num_hidden // 2, 3, padding=1)
+        self._adj1 = nn.Conv2d(self._num_hidden // 2, self._num_hidden, 3, padding=1)
+
+        self._conv2 = nn.Conv2d(self._num_hidden, self._num_hidden // 2, 3, padding=1)
+        self._adj2 = nn.Conv2d(self._num_hidden // 2, self._num_hidden, 3, padding=1)
+
+    def forward(self, batch: Tensor) -> Tensor:
+        batch = F.relu(self._stride(batch))
+        batch = batch + self._adj1(F.relu(self._conv1(batch)))
+        batch = batch + self._adj2(F.relu(self._conv2(batch)))
+        return batch
+
+    @property
+    def num_hidden(self):
+        return self._num_hidden
+
+
+# BN Complexity: 2
+class BNPreBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self._num_hidden = in_channels * 2
+        self._bn1 = nn.BatchNorm2d(in_channels)
+        self._conv1 = nn.Conv2d(in_channels, self._num_hidden, 3, stride=2, padding=1)
+        self._bn2 = nn.BatchNorm2d(self._num_hidden)
+        self._conv2 = nn.Conv2d(self._num_hidden, self._num_hidden - 1, 3, padding=1)
+        self._bn3 = nn.BatchNorm2d(self._num_hidden - 1)
+        self._conv3 = nn.Conv2d(self._num_hidden - 1, self._num_hidden, 3, padding=1)
+
+    def forward(self, batch: Tensor) -> Tensor:
+        batch = F.relu(self._conv1(self._bn1(batch)))
+        batch = F.relu(self._conv2(self._bn2(batch)))
+        batch = F.relu(self._conv3(self._bn3(batch)))
+        return batch
+
+    @property
+    def num_hidden(self):
+        return self._num_hidden
+
+
+class BNBetweenBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self._num_hidden = in_channels * 2
+        self._conv1 = nn.Conv2d(in_channels, self._num_hidden, 3, stride=2, padding=1)
+        self._bn1 = nn.BatchNorm2d(self._num_hidden)
+        self._conv2 = nn.Conv2d(self._num_hidden, self._num_hidden - 1, 3, padding=1)
+        self._bn2 = nn.BatchNorm2d(self._num_hidden - 1)
+        self._conv3 = nn.Conv2d(self._num_hidden - 1, self._num_hidden, 3, padding=1)
+        self._bn3 = nn.BatchNorm2d(self._num_hidden)
+
+    def forward(self, batch: Tensor) -> Tensor:
+        batch = F.relu(self._bn1(self._conv1(batch)))
+        batch = F.relu(self._bn2(self._conv2(batch)))
+        batch = F.relu(self._bn3(self._conv3(batch)))
+        return batch
+
+    @property
+    def num_hidden(self):
+        return self._num_hidden
+
+
+class BNPostBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self._num_hidden = in_channels * 2
+        self._conv1 = nn.Conv2d(in_channels, self._num_hidden, 3, stride=2, padding=1)
+        self._bn1 = nn.BatchNorm2d(self._num_hidden)
+        self._conv2 = nn.Conv2d(self._num_hidden, self._num_hidden - 1, 3, padding=1)
+        self._bn2 = nn.BatchNorm2d(self._num_hidden - 1)
+        self._conv3 = nn.Conv2d(self._num_hidden - 1, self._num_hidden, 3, padding=1)
+        self._bn3 = nn.BatchNorm2d(self._num_hidden)
+
+    def forward(self, batch: Tensor) -> Tensor:
+        batch = self._bn1(F.relu(self._conv1(batch)))
+        batch = self._bn2(F.relu(self._conv2(batch)))
+        batch = self._bn3(F.relu(self._conv3(batch)))
+        return batch
+
+    @property
+    def num_hidden(self):
+        return self._num_hidden
diff --git a/cnn_architecture/config.py b/cnn_architecture/config.py
@@ -0,0 +1,56 @@
+from pathlib import Path
+
+import numpy as np
+import torch
+from yacs.config import CfgNode as CN
+
+_cfg = CN()
+_cfg.NAME = ''
+_cfg.OUTPUT_DIR = ''
+_cfg.SEED = 42
+
+_cfg.MODEL = CN()
+_cfg.MODEL.CONV0 = CN()
+_cfg.MODEL.CONV0.IN_CHANNELS = 3
+_cfg.MODEL.CONV0.NUM_FILTERS = 128
+_cfg.MODEL.CONV0.SIZE = 7
+_cfg.MODEL.CONV0.STRIDE = 2
+
+_cfg.MODEL.NUM_CLASSES = 100
+_cfg.MODEL.BLOCK_TYPE = 'simple'
+
+_cfg.TRAIN = CN()
+_cfg.TRAIN.LR = 0.1
+_cfg.TRAIN.MOMENTUM = 0.9
+_cfg.TRAIN.WEIGHT_DECAY = 0.0005
+_cfg.TRAIN.NUM_ITERS = 10_000
+_cfg.TRAIN.BATCH_SIZE = 256
+_cfg.TRAIN.SHUFFLE = True
+
+_cfg.VAL = CN()
+_cfg.VAL.BATCH_SIZE = 0
+
+
+def load_cfg(cfg_path: Path):
+    cfg = _cfg.clone()
+    cfg.merge_from_file(cfg_path)
+    cfg = _transform_cfg(cfg, cfg_path.stem)
+
+    Path(cfg.OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+    _update_seeds(cfg.SEED)
+    return cfg
+
+
+def _transform_cfg(cfg: CN, name: str):
+    if cfg.NAME == '':
+        cfg.NAME = name
+    if cfg.OUTPUT_DIR == '':
+        cfg.OUTPUT_DIR = f"output/{cfg.NAME}"
+    if cfg.VAL.BATCH_SIZE == 0:
+        cfg.VAL.BATCH_SIZE = cfg.TRAIN.BATCH_SIZE * 2
+    return cfg
+
+
+def _update_seeds(seed: int):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
diff --git a/cnn_architecture/configs/add.yaml b/cnn_architecture/configs/add.yaml
@@ -0,0 +1,2 @@
+MODEL:
+  BLOCK_TYPE: 'add'
diff --git a/cnn_architecture/configs/bn_between.yaml b/cnn_architecture/configs/bn_between.yaml
@@ -0,0 +1,2 @@
+MODEL:
+  BLOCK_TYPE: 'bn_between'
diff --git a/cnn_architecture/configs/bn_post.yaml b/cnn_architecture/configs/bn_post.yaml
@@ -0,0 +1,2 @@
+MODEL:
+  BLOCK_TYPE: 'bn_post'
diff --git a/cnn_architecture/configs/bn_pre.yaml b/cnn_architecture/configs/bn_pre.yaml
@@ -0,0 +1,2 @@
+MODEL:
+  BLOCK_TYPE: 'bn_pre'
diff --git a/cnn_architecture/configs/simple.yaml b/cnn_architecture/configs/simple.yaml
diff --git a/cnn_architecture/model_builder.py b/cnn_architecture/model_builder.py
@@ -0,0 +1,37 @@
+import torch.nn.functional as F
+from torch import nn
+from yacs.config import CfgNode as CN
+
+from blocks import AddBlock, BNBetweenBlock, BNPostBlock, BNPreBlock, SimpleBlock
+
+
+class ModelTemplate(nn.Module):
+    def __init__(self, block_type, conv_cfg: CN, num_classes: int):
+        super().__init__()
+        self._conv = nn.Conv2d(conv_cfg.IN_CHANNELS, conv_cfg.NUM_FILTERS, conv_cfg.SIZE, stride=conv_cfg.STRIDE,
+                               padding=conv_cfg.SIZE // 2)
+
+        self._block1 = block_type(conv_cfg.NUM_FILTERS)
+        self._block2 = block_type(self._block1.num_hidden)
+        self._cls = nn.Linear(self._block2.num_hidden, num_classes)
+
+    def forward(self, batch):
+        batch = F.relu(self._conv(batch))
+        batch = self._block1(batch)
+        batch = self._block2(batch)
+        batch = F.max_pool2d(batch, batch.shape[-2:])
+        batch = self._cls(batch.view(batch.shape[:2]))
+        return batch
+
+
+_block_factory = {
+    'simple':     SimpleBlock,
+    'add':        AddBlock,
+    'bn_pre':     BNPreBlock,
+    'bn_between': BNBetweenBlock,
+    'bn_post':    BNPostBlock
+}
+
+
+def create_model(model_cfg: CN):
+    return ModelTemplate(_block_factory[model_cfg.BLOCK_TYPE], model_cfg.CONV0, model_cfg.NUM_CLASSES)
diff --git a/cnn_architecture/train_val.py b/cnn_architecture/train_val.py
@@ -0,0 +1,124 @@
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+import torch.optim as optim
+import torchvision
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from torchvision.transforms import Compose, Normalize, RandomCrop, RandomHorizontalFlip, ToTensor
+from tqdm import tqdm
+from yacs.config import CfgNode as CN
+
+from config import load_cfg
+from model_builder import create_model
+
+
+def train(model: nn.Module, train_cfg: CN, train_set: Dataset):
+    train_loader = DataLoader(train_set, batch_size=train_cfg.BATCH_SIZE, shuffle=train_cfg.SHUFFLE, num_workers=4)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=train_cfg.LR, momentum=train_cfg.MOMENTUM,
+                          weight_decay=train_cfg.WEIGHT_DECAY)
+
+    data_loader = iter(train_loader)
+
+    logging.info('Starting Training')
+    smooth_loss = 4.6  # Loss for random prediction of 100 classes
+    for idx in tqdm(range(train_cfg.NUM_ITERS), desc="Training"):
+        try:
+            batch, labels = next(data_loader)
+        except StopIteration:
+            data_loader = iter(train_loader)
+            batch, labels = next(data_loader)
+
+        batch = batch.cuda()
+        labels = labels.cuda()
+
+        outputs = model(batch)
+        loss = criterion(outputs, labels)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        smooth_loss = smooth_loss * 0.99 + float(loss) * 0.01
+
+        if (idx + 1) % 100 == 0:
+            logging.debug(f"{idx + 1:4d}: {smooth_loss:.3f} ({loss:.3f})")
+
+    logging.info('Finished Training')
+
+
+def val(model: nn.Module, val_cfg: CN, val_set: Dataset):
+    val_loader = torch.utils.data.DataLoader(val_set, batch_size=val_cfg.BATCH_SIZE, num_workers=4)
+
+    correct = 0
+    total = 0
+
+    class_correct = [0. for _ in range(100)]
+    class_total = [0. for _ in range(100)]
+    with torch.no_grad():
+        for images, labels in tqdm(val_loader, desc="Validating"):
+            images = images.cuda()
+            labels = labels.cuda()
+            outputs = model(images)
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+            c = (predicted == labels).squeeze()
+            for idx in range(labels.size(0)):
+                label = labels[idx]
+                class_correct[label] += c[idx].item()
+                class_total[label] += 1
+
+    logging.info(f"Accuracy of the network on the 10000 test images: {correct / total:.1%}")
+
+
+def main(cfg: CN):
+    logging.info(f"Using {cfg.NAME} configuration")
+
+    cache_path = Path(cfg.OUTPUT_DIR) / "final.model"
+    model = create_model(cfg.MODEL)
+    model = nn.DataParallel(model.cuda())
+
+    transform_list = [RandomCrop(32, padding=4), RandomHorizontalFlip(), ToTensor(),
+                      Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))]
+    if not cache_path.exists():
+        train_set = torchvision.datasets.CIFAR100(root='./data', train=True, download=True,
+                                                  transform=Compose(transform_list))
+        train(model, cfg.TRAIN, train_set)
+
+        torch.save(model.module.state_dict(), cache_path)
+
+    model.module.load_state_dict(torch.load(cache_path))
+
+    val_set = torchvision.datasets.CIFAR100(root='./data', train=False, download=True,
+                                            transform=Compose(transform_list[-2:]))
+    val(model, cfg.VAL, val_set)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('cfg', type=Path)
+    args = parser.parse_args()
+
+    # Ensure reproducibility
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+    _cfg = load_cfg(args.cfg)
+
+    logging
+    file_handler = logging.FileHandler(Path(_cfg.OUTPUT_DIR) / 'train_val.log')
+    file_handler.setLevel(logging.DEBUG)
+
+    stdout_handler = logging.StreamHandler()
+    stdout_handler.setLevel(logging.INFO)
+
+    logging.basicConfig(handlers=[file_handler, stdout_handler], level=logging.DEBUG,
+                        format="{asctime} [{levelname:^9s}] {message}", style="{",
+                        datefmt="%Y-%m-%d %H:%M:%S")
+
+    main(_cfg)

-Original file line number
+Diff line change
 data
 *.pkl
 *.onnx
 +*.model
 +*.log
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+MODEL:`
	`2`	`+ BLOCK_TYPE: 'bn_between'`