ENAS and DRATS search space zoo

examples/nas/darts/example_search_space.py → examples/nas/search_space_zoo/darts_example.py

@@ -9,13 +9,12 @@
 import torch.nn as nn
 
 import datasets
-# from model import CNN
 from nni.nas.pytorch.callbacks import ArchitectureCheckpoint, LRSchedulerCallback
 from nni.nas.pytorch.darts import DartsTrainer
 from utils import accuracy
 
 from nni.nas.pytorch.search_space_zoo import DartsCell
-from nni.nas.pytorch.search_space_zoo import DartsStackedCells
+from darts_search_space import DartsStackedCells
 
 logger = logging.getLogger('nni')
 

src/sdk/pynni/nni/nas/pytorch/search_space_zoo/darts_search_space.py → examples/nas/search_space_zoo/darts_stack_cells.py

@@ -22,20 +22,8 @@ class DartsStackedCells(nn.Module):
     n_layers: int
         the number of cells contained in this network
     factory_func: function
-        return a callable instance for demand cell structure
-        user should pass in ``__init__`` of the cell class with following parameters (see darts_cell for detail)
-            n_nodes: int
-                the number of nodes contained in this cell
-            channels_pp: int
-                the number of previous previous cell's output channels
-            channels_p: int
-                the number of previous cell's output channels
-            channels: int
-                the number of output channels for each node
-            reduction_p: bool
-                Is previous cell a reduction cell
-            reduction: bool
-                is current cell a reduction cell
+        return a callable instance for demand cell structure.
+        user should pass in ``__init__`` of the cell class with required parameters (see nni.nas.DartsCell for detail)
     n_nodes: int
         the number of nodes contained in each cell
     stem_multiplier: int

examples/nas/search_space_zoo/datasets.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import numpy as np
+import torch
+from torchvision import transforms
+from torchvision.datasets import CIFAR10
+
+
+class Cutout(object):
+    def __init__(self, length):
+        self.length = length
+
+    def __call__(self, img):
+        h, w = img.size(1), img.size(2)
+        mask = np.ones((h, w), np.float32)
+        y = np.random.randint(h)
+        x = np.random.randint(w)
+
+        y1 = np.clip(y - self.length // 2, 0, h)
+        y2 = np.clip(y + self.length // 2, 0, h)
+        x1 = np.clip(x - self.length // 2, 0, w)
+        x2 = np.clip(x + self.length // 2, 0, w)
+
+        mask[y1: y2, x1: x2] = 0.
+        mask = torch.from_numpy(mask)
+        mask = mask.expand_as(img)
+        img *= mask
+
+        return img
+
+
+def get_dataset(cls, cutout_length=0):
+    MEAN = [0.49139968, 0.48215827, 0.44653124]
+    STD = [0.24703233, 0.24348505, 0.26158768]
+    transf = [
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip()
+    ]
+    normalize = [
+        transforms.ToTensor(),
+        transforms.Normalize(MEAN, STD)
+    ]
+    cutout = []
+    if cutout_length > 0:
+        cutout.append(Cutout(cutout_length))
+
+    train_transform = transforms.Compose(transf + normalize + cutout)
+    valid_transform = transforms.Compose(normalize)
+
+    if cls == "cifar10":
+        dataset_train = CIFAR10(root="./data", train=True, download=True, transform=train_transform)
+        dataset_valid = CIFAR10(root="./data", train=False, download=True, transform=valid_transform)
+    else:
+        raise NotImplementedError
+    return dataset_train, dataset_valid

examples/nas/search_space_zoo/utils.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+def accuracy(output, target, topk=(1,)):
+    """ Computes the precision@k for the specified values of k """
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    # one-hot case
+    if target.ndimension() > 1:
+        target = target.max(1)[1]
+
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = dict()
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0)
+        res["acc{}".format(k)] = correct_k.mul_(1.0 / batch_size).item()
+    return res

src/sdk/pynni/nni/nas/pytorch/search_space_zoo/__init__.py

@@ -1,2 +1,2 @@
 from .darts_cell import DartsCell
-from .darts_search_space import DartsStackedCells
+