Merge pull request #3155 from microsoft/dev-retiarii

[Do NOT Squash] Merge retiarii dev branch to master

test/retiarii_test/darts/test.py

+import json
+import os
+import sys
+import torch
+from pathlib import Path
+
+from nni.retiarii.experiment import RetiariiExperiment, RetiariiExeConfig
+from nni.retiarii.strategies import TPEStrategy
+from nni.retiarii.trainer import PyTorchImageClassificationTrainer
+
+from darts_model import CNN
+
+if __name__ == '__main__':
+    base_model = CNN(32, 3, 16, 10, 8)
+    trainer = PyTorchImageClassificationTrainer(base_model, dataset_cls="CIFAR10",
+            dataset_kwargs={"root": "data/cifar10", "download": True},
+            dataloader_kwargs={"batch_size": 32},
+            optimizer_kwargs={"lr": 1e-3},
+            trainer_kwargs={"max_epochs": 1})
+
+    simple_startegy = TPEStrategy()
+
+    exp = RetiariiExperiment(base_model, trainer, [], simple_startegy)
+
+    exp_config = RetiariiExeConfig('local')
+    exp_config.experiment_name = 'darts_search'
+    exp_config.trial_concurrency = 2
+    exp_config.max_trial_number = 10
+    exp_config.trial_gpu_number = 1
+    exp_config.training_service.use_active_gpu = True
+    exp_config.training_service.gpu_indices = [1, 2]
+
+    exp.run(exp_config, 8081, debug=True)

test/retiarii_test/mnasnet/base_mnasnet.py

+import warnings
+
+import torch
+import torch.nn as torch_nn
+from torchvision.models.utils import load_state_dict_from_url
+import torch.nn.functional as F
+
+import sys
+from pathlib import Path
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+import nni.retiarii.nn.pytorch as nn
+from nni.retiarii import register_module
+
+# Paper suggests 0.9997 momentum, for TensorFlow. Equivalent PyTorch momentum is
+# 1.0 - tensorflow.
+_BN_MOMENTUM = 1 - 0.9997
+_FIRST_DEPTH = 32
+_MOBILENET_V2_FILTERS = [16, 24, 32, 64, 96, 160, 320]
+_MOBILENET_V2_NUM_LAYERS = [1, 2, 3, 4, 3, 3, 1]
+
+
+class _ResidualBlock(nn.Module):
+    def __init__(self, net):
+        super().__init__()
+        self.net = net
+
+    def forward(self, x):
+        return self.net(x) + x
+
+class _InvertedResidual(nn.Module):
+
+    def __init__(self, in_ch, out_ch, kernel_size, stride, expansion_factor, skip, bn_momentum=0.1):
+        super(_InvertedResidual, self).__init__()
+        assert stride in [1, 2]
+        assert kernel_size in [3, 5]
+        mid_ch = in_ch * expansion_factor
+        self.apply_residual = skip and in_ch == out_ch and stride == 1
+        self.layers = nn.Sequential(
+            # Pointwise
+            nn.Conv2d(in_ch, mid_ch, 1, bias=False),
+            nn.BatchNorm2d(mid_ch, momentum=bn_momentum),
+            nn.ReLU(inplace=True),
+            # Depthwise
+            nn.Conv2d(mid_ch, mid_ch, kernel_size, padding=kernel_size // 2,
+                      stride=stride, groups=mid_ch, bias=False),
+            nn.BatchNorm2d(mid_ch, momentum=bn_momentum),
+            nn.ReLU(inplace=True),
+            # Linear pointwise. Note that there's no activation.
+            nn.Conv2d(mid_ch, out_ch, 1, bias=False),
+            nn.BatchNorm2d(out_ch, momentum=bn_momentum))
+
+    def forward(self, input):
+        if self.apply_residual:
+            ret = self.layers(input) + input
+        else:
+            ret = self.layers(input)
+        return ret
+
+
+def _stack_inverted_residual(in_ch, out_ch, kernel_size, skip, stride, exp_factor, repeats, bn_momentum):
+    """ Creates a stack of inverted residuals. """
+    assert repeats >= 1
+    # First one has no skip, because feature map size changes.
+    first = _InvertedResidual(in_ch, out_ch, kernel_size, stride, exp_factor, skip, bn_momentum=bn_momentum)
+    remaining = []
+    for _ in range(1, repeats):
+        remaining.append(_InvertedResidual(out_ch, out_ch, kernel_size, 1, exp_factor, skip, bn_momentum=bn_momentum))
+    return nn.Sequential(first, *remaining)
+
+
+def _stack_normal_conv(in_ch, out_ch, kernel_size, skip, dconv, stride, repeats, bn_momentum):
+    assert repeats >= 1
+    stack = []
+    for i in range(repeats):
+        s = stride if i == 0 else 1
+        if dconv:
+            modules = [
+                nn.Conv2d(in_ch, in_ch, kernel_size, padding=kernel_size // 2, stride=s, groups=in_ch, bias=False),
+                nn.BatchNorm2d(in_ch, momentum=bn_momentum),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(in_ch, out_ch, 1, padding=0, stride=1, bias=False),
+                nn.BatchNorm2d(out_ch, momentum=bn_momentum)
+            ]
+        else:
+            modules = [
+                nn.Conv2d(in_ch, out_ch, kernel_size, padding=kernel_size // 2, stride=s, bias=False),
+                nn.ReLU(inplace=True),
+                nn.BatchNorm2d(out_ch, momentum=bn_momentum)
+            ]
+        if skip and in_ch == out_ch and s == 1:
+            # use different implementation for skip and noskip to align with pytorch
+            stack.append(_ResidualBlock(nn.Sequential(*modules)))
+        else:
+            stack += modules
+        in_ch = out_ch
+    return stack
+
+
+def _round_to_multiple_of(val, divisor, round_up_bias=0.9):
+    """ Asymmetric rounding to make `val` divisible by `divisor`. With default
+    bias, will round up, unless the number is no more than 10% greater than the
+    smaller divisible value, i.e. (83, 8) -> 80, but (84, 8) -> 88. """
+    assert 0.0 < round_up_bias < 1.0
+    new_val = max(divisor, int(val + divisor / 2) // divisor * divisor)
+    return new_val if new_val >= round_up_bias * val else new_val + divisor
+
+
+def _get_depths(depths, alpha):
+    """ Scales tensor depths as in reference MobileNet code, prefers rouding up
+    rather than down. """
+    return [_round_to_multiple_of(depth * alpha, 8) for depth in depths]
+
+@register_module()
+class MNASNet(nn.Module):
+    """ MNASNet, as described in https://arxiv.org/pdf/1807.11626.pdf. This
+    implements the B1 variant of the model.
+    >>> model = MNASNet(1000, 1.0)
+    >>> x = torch.rand(1, 3, 224, 224)
+    >>> y = model(x)
+    >>> y.dim()
+    1
+    >>> y.nelement()
+    1000
+    """
+    # Version 2 adds depth scaling in the initial stages of the network.
+    _version = 2
+
+    def __init__(self, alpha, depths, convops, kernel_sizes, num_layers,
+                 skips, num_classes=1000, dropout=0.2):
+        super(MNASNet, self).__init__()
+        assert alpha > 0.0
+        assert len(depths) == len(convops) == len(kernel_sizes) == len(num_layers) == len(skips) == 7
+        self.alpha = alpha
+        self.num_classes = num_classes
+        depths = _get_depths([_FIRST_DEPTH] + depths, alpha)
+        base_filter_sizes = [16, 24, 40, 80, 96, 192, 320]
+        exp_ratios = [3, 3, 3, 6, 6, 6, 6]
+        strides = [1, 2, 2, 2, 1, 2, 1]
+        layers = [
+            # First layer: regular conv.
+            nn.Conv2d(3, depths[0], 3, padding=1, stride=2, bias=False),
+            nn.BatchNorm2d(depths[0], momentum=_BN_MOMENTUM),
+            nn.ReLU(inplace=True),
+        ]
+        count = 0
+        #for conv, prev_depth, depth, ks, skip, stride, repeat, exp_ratio in \
+        #        zip(convops, depths[:-1], depths[1:], kernel_sizes, skips, strides, num_layers, exp_ratios):
+        for filter_size, exp_ratio, stride in zip(base_filter_sizes, exp_ratios, strides):
+            # TODO: restrict that "choose" can only be used within mutator
+            ph = nn.Placeholder(label=f'mutable_{count}', related_info={
+                           'kernel_size_options': [1, 3, 5],
+                           'n_layer_options': [1, 2, 3, 4],
+                           'op_type_options': ['__mutated__.base_mnasnet.RegularConv',
+                                               '__mutated__.base_mnasnet.DepthwiseConv',
+                                               '__mutated__.base_mnasnet.MobileConv'],
+                           #'se_ratio_options': [0, 0.25],
+                           'skip_options': ['identity', 'no'],
+                           'n_filter_options': [int(filter_size*x) for x in [0.75, 1.0, 1.25]],
+                           'exp_ratio': exp_ratio,
+                           'stride': stride,
+                           'in_ch': depths[0] if count == 0 else None
+            })
+            layers.append(ph)
+            '''if conv == "mconv":
+                # MNASNet blocks: stacks of inverted residuals.
+                layers.append(_stack_inverted_residual(prev_depth, depth, ks, skip,
+                                                       stride, exp_ratio, repeat, _BN_MOMENTUM))
+            else:
+                # Normal conv and depth-separated conv
+                layers += _stack_normal_conv(prev_depth, depth, ks, skip, conv == "dconv",
+                                             stride, repeat, _BN_MOMENTUM)'''
+            count += 1
+            if count >= 2:
+                break
+        layers += [
+            # Final mapping to classifier input.
+            nn.Conv2d(depths[7], 1280, 1, padding=0, stride=1, bias=False),
+            nn.BatchNorm2d(1280, momentum=_BN_MOMENTUM),
+            nn.ReLU(inplace=True),
+        ]
+        self.layers = nn.Sequential(*layers)
+        self.classifier = nn.Sequential(nn.Dropout(p=dropout, inplace=True),
+                                        nn.Linear(1280, num_classes))
+        self._initialize_weights()
+        #self.for_test = 10
+
+    def forward(self, x):
+        #if self.for_test == 10:
+        x = self.layers(x)
+        # Equivalent to global avgpool and removing H and W dimensions.
+        x = x.mean([2, 3])
+        x = F.relu(x)
+        return self.classifier(x)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch_nn.init.kaiming_normal_(m.weight, mode="fan_out",
+                                        nonlinearity="relu")
+                if m.bias is not None:
+                    torch_nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                torch_nn.init.ones_(m.weight)
+                torch_nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                torch_nn.init.kaiming_uniform_(m.weight, mode="fan_out",
+                                         nonlinearity="sigmoid")
+                torch_nn.init.zeros_(m.bias)
+
+
+def test_model(model):
+    model(torch.randn(2, 3, 224, 224))
+
+#====================definition of candidate op classes
+BN_MOMENTUM = 1 - 0.9997
+
+class RegularConv(nn.Module):
+    def __init__(self, kernel_size, in_ch, out_ch, skip, exp_ratio, stride):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.skip = skip
+        self.exp_ratio = exp_ratio
+        self.stride = stride
+
+        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=kernel_size // 2, stride=stride, bias=False)
+        self.relu = nn.ReLU(inplace=True)
+        self.bn = nn.BatchNorm2d(out_ch, momentum=BN_MOMENTUM)
+
+    def forward(self, x):
+        out = self.bn(self.relu(self.conv(x)))
+        if self.skip == 'identity':
+            out = out + x
+        return out
+
+class DepthwiseConv(nn.Module):
+    def __init__(self, kernel_size, in_ch, out_ch, skip, exp_ratio, stride):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.skip = skip
+        self.exp_ratio = exp_ratio
+        self.stride = stride
+
+        self.conv1 = nn.Conv2d(in_ch, in_ch, kernel_size, padding=kernel_size // 2, stride=stride, groups=in_ch, bias=False)
+        self.bn1 = nn.BatchNorm2d(in_ch, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(in_ch, out_ch, 1, padding=0, stride=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_ch, momentum=BN_MOMENTUM)
+
+    def forward(self, x):
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        if self.skip == 'identity':
+            out = out + x
+        return out
+
+class MobileConv(nn.Module):
+    def __init__(self, kernel_size, in_ch, out_ch, skip, exp_ratio, stride):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.skip = skip
+        self.exp_ratio = exp_ratio
+        self.stride = stride
+
+        mid_ch = in_ch * exp_ratio
+        self.layers = nn.Sequential(
+            # Pointwise
+            nn.Conv2d(in_ch, mid_ch, 1, bias=False),
+            nn.BatchNorm2d(mid_ch, momentum=BN_MOMENTUM),
+            nn.ReLU(inplace=True),
+            # Depthwise
+            nn.Conv2d(mid_ch, mid_ch, kernel_size, padding= (kernel_size - 1) // 2,
+                      stride=stride, groups=mid_ch, bias=False),
+            nn.BatchNorm2d(mid_ch, momentum=BN_MOMENTUM),
+            nn.ReLU(inplace=True),
+            # Linear pointwise. Note that there's no activation.
+            nn.Conv2d(mid_ch, out_ch, 1, bias=False),
+            nn.BatchNorm2d(out_ch, momentum=BN_MOMENTUM))
+
+    def forward(self, x):
+        out = self.layers(x)
+        if self.skip == 'identity':
+            out = out + x
+        return out
+
+# mnasnet0_5
+ir_module = _InvertedResidual(16, 16, 3, 1, 1, True)
\ No newline at end of file

test/retiarii_test/mnasnet/mutator.py

+import logging
+import sys
+from pathlib import Path
+sys.path.append(str(Path(__file__).resolve().parents[2]))
+from nni.retiarii import Mutator
+
+from base_mnasnet import RegularConv, DepthwiseConv, MobileConv
+
+_logger = logging.getLogger(__name__)
+
+class BlockMutator(Mutator):
+    def __init__(self, target: str):
+        super(BlockMutator, self).__init__()
+        self.target = target
+
+    def mutate(self, model):
+        nodes = model.get_nodes_by_label(self.target)
+        assert len(nodes) == 1
+        node = nodes[0]
+        graph = node.graph
+
+        related_info = node.operation.parameters
+        kernel_size = self.choice(related_info['kernel_size_options'])
+        op_type = self.choice(related_info['op_type_options'])
+        #self.choice(related_info['se_ratio_options'])
+        skip = self.choice(related_info['skip_options'])
+        n_filter = self.choice(related_info['n_filter_options'])
+
+        if related_info['in_ch'] is not None:
+            _logger.info('zql debug X ...')
+            in_ch = related_info['in_ch']
+        else:
+            assert len(node.predecessors) == 1
+            the_node = node.predecessors[0]
+            _logger.info('zql debug ...')
+            _logger.info(the_node.operation.parameters)
+            _logger.info(the_node.__repr__())
+            in_ch = the_node.operation.parameters['out_ch']
+
+        # update the placeholder to be a new operation
+        node.update_operation(op_type, {
+            'kernel_size': kernel_size,
+            'in_ch': in_ch,
+            'out_ch': n_filter,
+            'skip': 'no',
+            'exp_ratio': related_info['exp_ratio'],
+            'stride': related_info['stride']
+        })
+
+        # insert new nodes after the placeholder
+        n_layer = self.choice(related_info['n_layer_options'])
+        for i in range(1, n_layer):
+            node = graph.insert_node_on_edge(node.outgoing_edges[0],
+                                             '{}_{}'.format(self.target, i),
+                                             op_type,
+                                             {'kernel_size': kernel_size,
+                                              'in_ch': n_filter,
+                                              'out_ch': n_filter,
+                                              'skip': skip,
+                                              'exp_ratio': related_info['exp_ratio'],
+                                              'stride': 1})
+
+        # fix possible shape mismatch
+        # TODO: use formal method function to update parameters
+        if len(node.successors) == 1 and 'in_channels' in node.successors[0].operation.parameters:
+            node.successors[0].operation.parameters['in_channels'] = n_filter
\ No newline at end of file

test/retiarii_test/mnasnet/test.py

+import os
+import sys
+import torch
+from pathlib import Path
+
+from nni.retiarii.trainer import PyTorchImageClassificationTrainer
+
+from base_mnasnet import MNASNet
+from nni.retiarii.experiment import RetiariiExperiment, RetiariiExeConfig
+
+from nni.retiarii.strategies import TPEStrategy
+from mutator import BlockMutator
+
+if __name__ == '__main__':
+    _DEFAULT_DEPTHS = [16, 24, 40, 80, 96, 192, 320]
+    _DEFAULT_CONVOPS = ["dconv", "mconv", "mconv", "mconv", "mconv", "mconv", "mconv"]
+    _DEFAULT_SKIPS = [False, True, True, True, True, True, True]
+    _DEFAULT_KERNEL_SIZES = [3, 3, 5, 5, 3, 5, 3]
+    _DEFAULT_NUM_LAYERS = [1, 3, 3, 3, 2, 4, 1]
+
+    base_model = MNASNet(0.5, _DEFAULT_DEPTHS, _DEFAULT_CONVOPS, _DEFAULT_KERNEL_SIZES,
+                    _DEFAULT_NUM_LAYERS, _DEFAULT_SKIPS)
+    trainer = PyTorchImageClassificationTrainer(base_model, dataset_cls="CIFAR10",
+            dataset_kwargs={"root": "data/cifar10", "download": True},
+            dataloader_kwargs={"batch_size": 32},
+            optimizer_kwargs={"lr": 1e-3},
+            trainer_kwargs={"max_epochs": 1})
+
+    # new interface
+    applied_mutators = []
+    applied_mutators.append(BlockMutator('mutable_0'))
+    applied_mutators.append(BlockMutator('mutable_1'))
+
+    simple_startegy = TPEStrategy()
+
+    exp = RetiariiExperiment(base_model, trainer, applied_mutators, simple_startegy)
+
+    exp_config = RetiariiExeConfig('local')
+    exp_config.experiment_name = 'mnasnet_search'
+    exp_config.trial_concurrency = 2
+    exp_config.max_trial_number = 10
+    exp_config.training_service.use_active_gpu = False
+
+    exp.run(exp_config, 8081, debug=True)

test/retiarii_test/simple_strategy.py

+import json
+import logging
+import random
+import os
+
+from nni.retiarii import Model, submit_models, wait_models
+from nni.retiarii.strategy import BaseStrategy
+from nni.retiarii import Sampler
+
+
+_logger = logging.getLogger(__name__)
+
+class RandomSampler(Sampler):
+    def choice(self, candidates, mutator, model, index):
+        return random.choice(candidates)
+
+class SimpleStrategy(BaseStrategy):
+    def __init__(self):
+        self.name = ''
+
+    def run(self, base_model, applied_mutators, trainer):
+        try:
+            _logger.info('stargety start...')
+            while True:
+                model = base_model
+                _logger.info('apply mutators...')
+                _logger.info('mutators: {}'.format(applied_mutators))
+                random_sampler = RandomSampler()
+                for mutator in applied_mutators:
+                    _logger.info('mutate model...')
+                    mutator.bind_sampler(random_sampler)
+                    model = mutator.apply(model)
+                # get and apply training approach
+                _logger.info('apply training approach...')
+                model.apply_trainer(trainer['modulename'], trainer['args'])
+                # run models
+                submit_models(model)
+                wait_models(model)
+                _logger.info('Strategy says:', model.metric)
+        except Exception as e:
+            _logger.error(logging.exception('message'))

test/ut/__init__.py

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
 """
 Unit test of NNI Python modules.
 
 Test cases of each module should be placed at same path of their source files.
-For example if `nni/tool/annotation` has one test case, it should be placed at `test/ut/tool/annotation.py`;
+For example if `nni/tool/annotation` has one test case, it should be placed at `test/ut/tool/test_annotation.py`;
 if it has multiple test cases, they should be placed in `test/ut/tool/annotation/` directory.
 
 "Legacy" test cases carried from NNI v1.x might not follow above convention:
 
   + Directory `sdk` contains old test cases previously in `src/sdk/pynni/tests`.
-  + Directory `tools/cmd` contains old test cases previously in `tools/cmd/tests`.
+  + Directory `tools/nnictl` contains old test cases previously in `tools/nni_cmd/tests`.
   + Directory `tools/annotation` contains old test cases previously in `tools/nni_annotation`.
   + Directory `tools/trial_tool` contains old test cases previously in `tools/nni_trial_tool/test`.
 """
+
+import os
+
+os.environ['NNI_PLATFORM'] = 'unittest'
+os.environ['NNI_TRIAL_JOB_ID'] = 'test_trial_job_id'

test/ut/retiarii/converted_mnist_pytorch.json

+{
+    "_model__stem":{
+       "inputs":[
+          "_inputs__1"
+       ],
+       "outputs":[
+          "pool2__1"
+       ],
+       "nodes":{
+          "_model__stem__conv1":{
+             "operation":{
+                "type":"__torch__.torch.nn.modules.conv.Conv2d",
+                "parameters":{
+                   "out_channels":32,
+                   "in_channels":1,
+                   "kernel_size":5
+                }
+             }
+          },
+          "_model__stem__pool1":{
+             "operation":{
+                "type":"__torch__.torch.nn.modules.pooling.MaxPool2d",
+                "parameters":{
+                   "kernel_size":2
+                }
+             }
+          },
+          "_model__stem__conv2":{
+             "operation":{
+                "type":"__torch__.torch.nn.modules.conv.Conv2d",
+                "parameters":{
+                   "out_channels":64,
+                   "in_channels":32,
+                   "kernel_size":5
+                }
+             }
+          },
+          "_model__stem__pool2":{
+             "operation":{
+                "type":"__torch__.torch.nn.modules.pooling.MaxPool2d",
+                "parameters":{
+                   "kernel_size":2
+                }
+             }
+          }
+       },
+       "edges":[
+          {
+             "head":[
+                "_inputs",
+                0
+             ],
+             "tail":[
+                "_model__stem__conv1",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__stem__conv1",
+                null
+             ],
+             "tail":[
+                "_model__stem__pool1",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__stem__pool1",
+                null
+             ],
+             "tail":[
+                "_model__stem__conv2",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__stem__conv2",
+                null
+             ],
+             "tail":[
+                "_model__stem__pool2",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__stem__pool2",
+                null
+             ],
+             "tail":[
+                "_outputs",
+                null
+             ]
+          }
+       ]
+    },
+    "_model":{
+       "inputs":[
+          "image__1"
+       ],
+       "outputs":[
+          "softmax__1"
+       ],
+       "nodes":{
+          "_model__Constant2":{
+             "operation":{
+                "type":"prim::Constant",
+                "parameters":{
+                   
+                }
+             }
+          },
+          "_model__Constant3":{
+             "operation":{
+                "type":"prim::Constant",
+                "parameters":{
+                   "value":3
+                }
+             }
+          },
+          "_model__Constant4":{
+             "operation":{
+                "type":"prim::Constant",
+                "parameters":{
+                   "value":-1
+                }
+             }
+          },
+          "_model__Constant5":{
+             "operation":{
+                "type":"prim::Constant",
+                "parameters":{
+                   "value":0
+                }
+             }
+          },
+          "_model__stem":{
+             "operation":{
+                "type":"_cell",
+                "parameters":{
+                   
+                },
+                "cell_name":"_model__stem"
+             }
+          },
+          "_model__Size6":{
+             "operation":{
+                "type":"aten::size",
+                "parameters":{
+                   
+                }
+             }
+          },
+          "_model__ListConstruct7":{
+             "operation":{
+                "type":"prim::ListConstruct",
+                "parameters":{
+                   
+                }
+             }
+          },
+          "_model__View8":{
+             "operation":{
+                "type":"aten::view",
+                "parameters":{
+                   
+                }
+             }
+          },
+          "_model__fc1":{
+             "operation":{
+                "type":"__torch__.torch.nn.modules.linear.Linear",
+                "parameters":{
+                   "in_features":1024,
+                   "out_features":256
+                }
+             }
+          },
+          "_model__fc2":{
+             "operation":{
+                "type":"__torch__.torch.nn.modules.linear.Linear",
+                "parameters":{
+                   "in_features":256,
+                   "out_features":10
+                }
+             }
+          },
+          "_model__softmax9":{
+             "operation":{
+                "type":"Function.softmax",
+                "parameters":{
+                   
+                }
+             }
+          }
+       },
+       "edges":[
+          {
+             "head":[
+                "_inputs",
+                0
+             ],
+             "tail":[
+                "_model__stem",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__stem",
+                null
+             ],
+             "tail":[
+                "_model__Size6",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__Constant5",
+                null
+             ],
+             "tail":[
+                "_model__Size6",
+                1
+             ]
+          },
+          {
+             "head":[
+                "_model__Size6",
+                null
+             ],
+             "tail":[
+                "_model__ListConstruct7",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__Constant4",
+                null
+             ],
+             "tail":[
+                "_model__ListConstruct7",
+                1
+             ]
+          },
+          {
+             "head":[
+                "_model__stem",
+                null
+             ],
+             "tail":[
+                "_model__View8",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__ListConstruct7",
+                null
+             ],
+             "tail":[
+                "_model__View8",
+                1
+             ]
+          },
+          {
+             "head":[
+                "_model__View8",
+                null
+             ],
+             "tail":[
+                "_model__fc1",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__fc1",
+                null
+             ],
+             "tail":[
+                "_model__fc2",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__fc2",
+                null
+             ],
+             "tail":[
+                "_model__softmax9",
+                0
+             ]
+          },
+          {
+             "head":[
+                "_model__Constant4",
+                null
+             ],
+             "tail":[
+                "_model__softmax9",
+                1
+             ]
+          },
+          {
+             "head":[
+                "_model__Constant3",
+                null
+             ],
+             "tail":[
+                "_model__softmax9",
+                2
+             ]
+          },
+          {
+             "head":[
+                "_model__Constant2",
+                null
+             ],
+             "tail":[
+                "_model__softmax9",
+                3
+             ]
+          },
+          {
+             "head":[
+                "_model__softmax9",
+                null
+             ],
+             "tail":[
+                "_outputs",
+                null
+             ]
+          }
+       ]
+    },
+    "_training_config": {
+        "module": "nni.retiarii.trainer.PyTorchImageClassificationTrainer",
+        "kwargs": {
+            "dataset_cls": "MNIST",
+            "dataset_kwargs": {
+                "root": "data/mnist",
+                "download": true
+            },
+            "dataloader_kwargs": {
+                "batch_size": 32
+            },
+            "optimizer_cls" : "SGD",
+            "optimizer_kwargs": {
+                "lr": 1e-3
+            },
+            "trainer_kwargs": {
+                "max_epochs": 1
+            }
+        }
+    }
+ }
\ No newline at end of file

test/ut/retiarii/debug_mnist_pytorch.py

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+
+class _model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.stem = stem()
+        
+        self.fc1 = nn.Linear(1024, 256)
+        self.fc2 = nn.Linear(256, 10)
+        
+
+    def forward(self, image):
+        stem = self.stem(image)
+        flatten = stem.view(stem.size(0), -1)
+        fc1 = self.fc1(flatten)
+        fc2 = self.fc2(fc1)
+        softmax = F.softmax(fc2, -1)
+        return softmax
+
+
+
+class stem(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(out_channels=32, in_channels=1, kernel_size=5)
+        self.pool1 = nn.MaxPool2d(kernel_size=2)
+        self.conv2 = nn.Conv2d(out_channels=64, in_channels=32, kernel_size=5)
+        self.pool2 = nn.MaxPool2d(kernel_size=2)
+
+    def forward(self, *_inputs):
+        conv1 = self.conv1(_inputs[0])
+        pool1 = self.pool1(conv1)
+        conv2 = self.conv2(pool1)
+        pool2 = self.pool2(conv2)
+        return pool2

test/ut/retiarii/dedup_logical_graph.json

+{"inputs": null, "outputs": null, "nodes": {"2__outputs": {"operation": {"type": "_outputs", "parameters": {}}}, "2__model__Constant2": {"operation": {"type": "prim::Constant", "parameters": {}}}, "2__model__Constant3": {"operation": {"type": "prim::Constant", "parameters": {"value": 3}}}, "2__model__Constant4": {"operation": {"type": "prim::Constant", "parameters": {"value": -1}}}, "2__model__Constant5": {"operation": {"type": "prim::Constant", "parameters": {"value": 0}}}, "2__model__stem": {"operation": {"type": "_cell", "parameters": {}, "cell_name": "_model__stem"}}, "2__model__Size6": {"operation": {"type": "aten::size", "parameters": {}}}, "2__model__ListConstruct7": {"operation": {"type": "prim::ListConstruct", "parameters": {}}}, "2__model__View8": {"operation": {"type": "aten::view", "parameters": {}}}, "2__model__fc1": {"operation": {"type": "__torch__.torch.nn.modules.linear.Linear", "parameters": {"in_features": 1024, "out_features": 256}}}, "2__model__fc2": {"operation": {"type": "__torch__.torch.nn.modules.linear.Linear", "parameters": {"in_features": 256, "out_features": 10}}}, "2__model__softmax9": {"operation": {"type": "Function.softmax", "parameters": {}}}, "3__outputs": {"operation": {"type": "_outputs", "parameters": {}}}, "3__model__Constant2": {"operation": {"type": "prim::Constant", "parameters": {}}}, "3__model__Constant3": {"operation": {"type": "prim::Constant", "parameters": {"value": 3}}}, "3__model__Constant4": {"operation": {"type": "prim::Constant", "parameters": {"value": -1}}}, "3__model__Constant5": {"operation": {"type": "prim::Constant", "parameters": {"value": 0}}}, "3__model__stem": {"operation": {"type": "_cell", "parameters": {}, "cell_name": "_model__stem"}}, "3__model__Size6": {"operation": {"type": "aten::size", "parameters": {}}}, "3__model__ListConstruct7": {"operation": {"type": "prim::ListConstruct", "parameters": {}}}, "3__model__View8": {"operation": {"type": "aten::view", "parameters": {}}}, "3__model__fc1": {"operation": {"type": "__torch__.torch.nn.modules.linear.Linear", "parameters": {"in_features": 1024, "out_features": 256}}}, "3__model__fc2": {"operation": {"type": "__torch__.torch.nn.modules.linear.Linear", "parameters": {"in_features": 256, "out_features": 10}}}, "3__model__softmax9": {"operation": {"type": "Function.softmax", "parameters": {}}}, "4__outputs": {"operation": {"type": "_outputs", "parameters": {}}}, "4__model__Constant2": {"operation": {"type": "prim::Constant", "parameters": {}}}, "4__model__Constant3": {"operation": {"type": "prim::Constant", "parameters": {"value": 3}}}, "4__model__Constant4": {"operation": {"type": "prim::Constant", "parameters": {"value": -1}}}, "4__model__Constant5": {"operation": {"type": "prim::Constant", "parameters": {"value": 0}}}, "4__model__stem": {"operation": {"type": "_cell", "parameters": {}, "cell_name": "_model__stem"}}, "4__model__Size6": {"operation": {"type": "aten::size", "parameters": {}}}, "4__model__ListConstruct7": {"operation": {"type": "prim::ListConstruct", "parameters": {}}}, "4__model__View8": {"operation": {"type": "aten::view", "parameters": {}}}, "4__model__fc1": {"operation": {"type": "__torch__.torch.nn.modules.linear.Linear", "parameters": {"in_features": 1024, "out_features": 256}}}, "4__model__fc2": {"operation": {"type": "__torch__.torch.nn.modules.linear.Linear", "parameters": {"in_features": 256, "out_features": 10}}}, "4__model__softmax9": {"operation": {"type": "Function.softmax", "parameters": {}}}, "1_Dedup__inputs": {"operation": {"type": "_inputs", "parameters": {}}}}, "edges": [["Dedup__inputs", "2__model__stem"], ["2__model__stem", "2__model__Size6"], ["2__model__Constant5", "2__model__Size6"], ["2__model__Size6", "2__model__ListConstruct7"], ["2__model__Constant4", "2__model__ListConstruct7"], ["2__model__stem", "2__model__View8"], ["2__model__ListConstruct7", "2__model__View8"], ["2__model__View8", "2__model__fc1"], ["2__model__fc1", "2__model__fc2"], ["2__model__fc2", "2__model__softmax9"], ["2__model__Constant4", "2__model__softmax9"], ["2__model__Constant3", "2__model__softmax9"], ["2__model__Constant2", "2__model__softmax9"], ["2__model__softmax9", "2__outputs"], ["Dedup__inputs", "3__model__stem"], ["3__model__stem", "3__model__Size6"], ["3__model__Constant5", "3__model__Size6"], ["3__model__Size6", "3__model__ListConstruct7"], ["3__model__Constant4", "3__model__ListConstruct7"], ["3__model__stem", "3__model__View8"], ["3__model__ListConstruct7", "3__model__View8"], ["3__model__View8", "3__model__fc1"], ["3__model__fc1", "3__model__fc2"], ["3__model__fc2", "3__model__softmax9"], ["3__model__Constant4", "3__model__softmax9"], ["3__model__Constant3", "3__model__softmax9"], ["3__model__Constant2", "3__model__softmax9"], ["3__model__softmax9", "3__outputs"], ["Dedup__inputs", "4__model__stem"], ["4__model__stem", "4__model__Size6"], ["4__model__Constant5", "4__model__Size6"], ["4__model__Size6", "4__model__ListConstruct7"], ["4__model__Constant4", "4__model__ListConstruct7"], ["4__model__stem", "4__model__View8"], ["4__model__ListConstruct7", "4__model__View8"], ["4__model__View8", "4__model__fc1"], ["4__model__fc1", "4__model__fc2"], ["4__model__fc2", "4__model__softmax9"], ["4__model__Constant4", "4__model__softmax9"], ["4__model__Constant3", "4__model__softmax9"], ["4__model__Constant2", "4__model__softmax9"], ["4__model__softmax9", "4__outputs"]]}
\ No newline at end of file

test/ut/retiarii/mnist-tensorflow.json

+{
+    "_model": {
+        "inputs": ["image"],
+        "outputs": ["metric"],
+
+        "nodes": {
+            "stem": {"operation": {"type": "_cell", "parameters": {}, "cell_name": "stem"}},
+            "flatten": {"operation": {"type": "Flatten", "parameters": {}}},
+            "fc1": {"operation": {"type": "Dense", "parameters": {"units": 1024, "activation": "relu"}}},
+            "fc2": {"operation": {"type": "Dense", "parameters": {"units": 10}}},
+            "softmax": {"operation": {"type": "Softmax", "parameters": {}}}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["stem", 0]},
+            {"head": ["stem", 0], "tail": ["flatten", null]},
+            {"head": ["flatten", null], "tail": ["fc1", null]},
+            {"head": ["fc1", null], "tail": ["fc2", null]},
+            {"head": ["fc2", null], "tail": ["softmax", null]},
+            {"head": ["softmax", null], "tail": ["_outputs", 0]}
+        ]
+    },
+
+    "stem": {
+        "nodes": {
+            "conv1": {"operation": {"type": "Conv2D", "parameters": {"filters": 32, "kernel_size": 5, "activation": "relu"}}},
+            "pool1": {"operation": {"type": "MaxPool2D", "parameters": {"pool_size": 2}}},
+            "conv2": {"operation": {"type": "Conv2D", "parameters": {"filters": 64, "kernel_size": 5, "activation": "relu"}}},
+            "pool2": {"operation": {"type": "MaxPool2D", "parameters": {"pool_size": 2}}}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["conv1", null]},
+            {"head": ["conv1", null], "tail": ["pool1", null]},
+            {"head": ["pool1", null], "tail": ["conv2", null]},
+            {"head": ["conv2", null], "tail": ["pool2", null]},
+            {"head": ["pool2", null], "tail": ["_outputs", 0]}
+        ]
+    },
+
+    "_training_config": {
+        "module": "_debug_no_trainer",
+        "kwargs": {}
+    }
+}

test/ut/retiarii/mnist_pytorch.json

+{
+    "_model": {
+        "inputs": ["image"],
+        "outputs": ["metric"],
+
+        "nodes": {
+            "stem": {"operation": {"type": "_cell", "cell_name": "stem"}},
+            "flatten": {"operation": {"type": "Flatten"}},
+            "fc1": {"operation": {"type": "Dense", "parameters": {"out_features": 256, "in_features": 1024}}},
+            "fc2": {"operation": {"type": "Dense", "parameters": {"out_features": 10, "in_features": 256}}},
+            "softmax": {"operation": {"type": "Softmax"}}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["stem", null]},
+            {"head": ["stem", null], "tail": ["flatten", null]},
+            {"head": ["flatten", null], "tail": ["fc1", null]},
+            {"head": ["fc1", null], "tail": ["fc2", null]},
+            {"head": ["fc2", null], "tail": ["softmax", null]},
+            {"head": ["softmax", null], "tail": ["_outputs", 0]}
+        ]
+    },
+
+    "stem": {
+        "nodes": {
+            "conv1": {"operation": {"type": "__torch__.Conv2d", "parameters": {"out_channels": 32, "in_channels": 1, "kernel_size": 5}}},
+            "pool1": {"operation": {"type": "__torch__.MaxPool2d", "parameters": {"kernel_size": 2}}},
+            "conv2": {"operation": {"type": "__torch__.Conv2d", "parameters": {"out_channels": 64, "in_channels": 32, "kernel_size": 5}}},
+            "pool2": {"operation": {"type": "__torch__.MaxPool2d", "parameters": {"kernel_size": 2}}}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["conv1", null]},
+            {"head": ["conv1", null], "tail": ["pool1", null]},
+            {"head": ["pool1", null], "tail": ["conv2", null]},
+            {"head": ["conv2", null], "tail": ["pool2", null]},
+            {"head": ["pool2", null], "tail": ["_outputs", 0]}
+        ]
+    },
+
+    "_training_config": {
+        "module": "nni.retiarii.trainer.PyTorchImageClassificationTrainer",
+        "kwargs": {
+            "dataset_cls": "MNIST",
+            "dataset_kwargs": {
+                "root": "data/mnist",
+                "download": true
+            },
+            "dataloader_kwargs": {
+                "batch_size": 32
+            },
+            "optimizer_cls" : "SGD",
+            "optimizer_kwargs": {
+                "lr": 1e-3
+            },
+            "trainer_kwargs": {
+                "max_epochs": 1
+            }
+        }
+    }
+}

test/ut/retiarii/test_cgo_engine.py

+import json
+import os
+import sys
+import threading
+import unittest
+import logging
+import time
+import torch
+
+from pathlib import Path
+
+from nni.retiarii.execution.cgo_engine import CGOExecutionEngine
+from nni.retiarii.execution.logical_optimizer.logical_plan import LogicalPlan
+from nni.retiarii.execution.logical_optimizer.opt_dedup_input import DedupInputOptimizer
+from nni.retiarii.codegen import model_to_pytorch_script
+from nni.retiarii import Model, Node
+
+from nni.retiarii import Model, submit_models
+from nni.retiarii.codegen import model_to_pytorch_script
+from nni.retiarii.integration import RetiariiAdvisor
+from nni.retiarii.trainer import PyTorchImageClassificationTrainer, PyTorchMultiModelTrainer
+from nni.retiarii.utils import import_
+
+
+def _load_mnist(n_models: int = 1):
+    path = Path(__file__).parent / 'converted_mnist_pytorch.json'
+    with open(path) as f:
+        mnist_model = Model._load(json.load(f))
+    if n_models == 1:
+        return mnist_model
+    else:
+        models = [mnist_model]
+        for i in range(n_models-1):
+            models.append(mnist_model.fork())
+        return models
+
+
+@unittest.skip('Skipped in this version')
+class CGOEngineTest(unittest.TestCase):
+
+    def test_submit_models(self):
+        os.environ['CGO'] = 'true'
+        os.makedirs('generated', exist_ok=True)
+        from nni.runtime import protocol, platform
+        import nni.runtime.platform.test as tt
+        protocol._out_file = open('generated/debug_protocol_out_file.py', 'wb')
+        protocol._in_file = open('generated/debug_protocol_out_file.py', 'rb')
+
+        models = _load_mnist(2)
+        advisor = RetiariiAdvisor()
+        submit_models(*models)
+
+        if torch.cuda.is_available() and torch.cuda.device_count() >= 2:
+            cmd, data = protocol.receive()
+            params = json.loads(data)
+            params['parameters']['training_kwargs']['max_steps'] = 100
+
+            tt.init_params(params)
+
+            trial_thread = threading.Thread(target=CGOExecutionEngine.trial_execute_graph())
+            trial_thread.start()
+            last_metric = None
+            while True:
+                time.sleep(1)
+                if tt._last_metric:
+                    metric = tt.get_last_metric()
+                    if metric == last_metric:
+                        continue
+                    advisor.handle_report_metric_data(metric)
+                    last_metric = metric
+                if not trial_thread.is_alive():
+                    break
+
+            trial_thread.join()
+        advisor.stopping = True
+        advisor.default_worker.join()
+        advisor.assessor_worker.join()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/ut/retiarii/test_dedup_input.py

+import json
+import os
+import sys
+import threading
+import unittest
+import logging
+import time
+
+from pathlib import Path
+
+from nni.retiarii.execution.cgo_engine import CGOExecutionEngine
+from nni.retiarii.execution.logical_optimizer.logical_plan import LogicalPlan
+from nni.retiarii.execution.logical_optimizer.opt_dedup_input import DedupInputOptimizer
+from nni.retiarii.codegen import model_to_pytorch_script
+from nni.retiarii import Model, Node
+
+from nni.retiarii import Model, submit_models
+from nni.retiarii.codegen import model_to_pytorch_script
+from nni.retiarii.integration import RetiariiAdvisor
+from nni.retiarii.trainer import PyTorchImageClassificationTrainer, PyTorchMultiModelTrainer
+from nni.retiarii.utils import import_
+
+
+def _load_mnist(n_models: int = 1):
+    path = Path(__file__).parent / 'converted_mnist_pytorch.json'
+    with open(path) as f:
+        mnist_model = Model._load(json.load(f))
+    if n_models == 1:
+        return mnist_model
+    else:
+        models = [mnist_model]
+        for i in range(n_models-1):
+            models.append(mnist_model.fork())
+        return models
+
+
+@unittest.skip('Skipped in this version')
+class DedupInputTest(unittest.TestCase):
+    def _build_logical_with_mnist(self, n_models: int):
+        lp = LogicalPlan()
+        models = _load_mnist(n_models=n_models)
+        for m in models:
+            lp.add_model(m)
+        return lp, models
+
+    def _test_add_model(self):
+        lp, models = self._build_logical_with_mnist(3)
+        for node in lp.logical_graph.hidden_nodes:
+            old_nodes = [m.root_graph.get_node_by_id(node.id) for m in models]
+
+            self.assertTrue(any([old_nodes[0].__repr__() == Node.__repr__(x) for x in old_nodes]))
+
+    def test_dedup_input(self):
+        os.environ['CGO'] = 'true'
+        lp, models = self._build_logical_with_mnist(3)
+        opt = DedupInputOptimizer()
+        opt.convert(lp)
+        with open('dedup_logical_graph.json', 'r') as fp:
+            correct_dump = fp.readlines()
+        lp_dump = lp.logical_graph._dump()
+
+        self.assertTrue(correct_dump[0] == json.dumps(lp_dump))
+
+        advisor = RetiariiAdvisor()
+        cgo = CGOExecutionEngine()
+
+        phy_models = cgo._assemble(lp)
+        self.assertTrue(len(phy_models) == 1)
+        # logging.info(phy_models[0][0]._dump())
+        # script=model_to_pytorch_script(phy_models[0][0], placement = phy_models[0][1])
+        # logging.info(script)
+        # with open('generated/debug_dedup_input.py', 'w') as fp:
+        #     fp.write(script)
+        # sys.path.insert(0, 'generated')
+        # multi_model = import_('debug_dedup_input.logical_0')
+        # trainer = PyTorchMultiModelTrainer(
+        #     multi_model(), phy_models[0][0].training_config.kwargs
+        # )
+        # trainer.fit()
+
+        advisor.stopping = True
+        advisor.default_worker.join()
+        advisor.assessor_worker.join()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/ut/retiarii/test_engine.py

+import json
+import os
+import sys
+import threading
+import unittest
+from pathlib import Path
+
+import nni
+from nni.retiarii import Model, submit_models
+from nni.retiarii.codegen import model_to_pytorch_script
+from nni.retiarii.integration import RetiariiAdvisor, register_advisor
+from nni.retiarii.trainer import PyTorchImageClassificationTrainer
+from nni.retiarii.utils import import_
+
+
+@unittest.skip('Skipped in this version')
+class CodeGenTest(unittest.TestCase):
+    def test_mnist_example_pytorch(self):
+        with open('mnist_pytorch.json') as f:
+            model = Model._load(json.load(f))
+            script = model_to_pytorch_script(model)
+        with open('debug_mnist_pytorch.py') as f:
+            reference_script = f.read()
+        self.assertEqual(script.strip(), reference_script.strip())
+
+
+@unittest.skip('Skipped in this version')
+class TrainerTest(unittest.TestCase):
+    def test_trainer(self):
+        sys.path.insert(0, Path(__file__).parent.as_posix())
+        Model = import_('debug_mnist_pytorch._model')
+        trainer = PyTorchImageClassificationTrainer(
+            Model(),
+            dataset_kwargs={'root': (Path(__file__).parent / 'data' / 'mnist').as_posix(), 'download': True},
+            dataloader_kwargs={'batch_size': 32},
+            optimizer_kwargs={'lr': 1e-3},
+            trainer_kwargs={'max_epochs': 1}
+        )
+        trainer.fit()
+
+
+@unittest.skip('Skipped in this version')
+class EngineTest(unittest.TestCase):
+
+    def test_submit_models(self):
+        os.makedirs('generated', exist_ok=True)
+        from nni.runtime import protocol
+        protocol._out_file = open(Path(__file__).parent / 'generated/debug_protocol_out_file.py', 'wb')
+        advisor = RetiariiAdvisor()
+        with open('mnist_pytorch.json') as f:
+            model = Model._load(json.load(f))
+        submit_models(model, model)
+
+        advisor.stopping = True
+        advisor.default_worker.join()
+        advisor.assessor_worker.join()
+
+    def test_execution_engine(self):
+        pass

test/ut/retiarii/test_graph.py

+import json
+from pathlib import Path
+import sys
+
+from nni.retiarii import *
+
+
+json_files = [
+    'mnist-tensorflow.json'
+]
+
+
+def test_model_load_dump():
+    for json_file in json_files:
+        path = Path(__file__).parent / json_file
+        _test_file(path)
+
+
+def _test_file(json_path):
+    orig_ir = json.load(json_path.open())
+    model = Model._load(orig_ir)
+    dump_ir = model._dump()
+
+    # add default values to JSON, so we can compare with `==`
+    for graph_name, graph in orig_ir.items():
+        if graph_name == '_training_config':
+            continue
+        if 'inputs' not in graph:
+            graph['inputs'] = None
+        if 'outputs' not in graph:
+            graph['outputs'] = None
+
+    # debug output
+    #json.dump(orig_ir, open('_orig.json', 'w'), indent=4)
+    #json.dump(dump_ir, open('_dump.json', 'w'), indent=4)
+    assert orig_ir == dump_ir
+
+
+if __name__ == '__main__':
+    test_model_load_dump()

test/ut/retiarii/test_mutator.py

+import json
+from pathlib import Path
+import sys
+
+from nni.retiarii import *
+
+# FIXME
+import nni.retiarii.debug_configs
+nni.retiarii.debug_configs.framework = 'tensorflow'
+
+max_pool = Operation.new('MaxPool2D', {'pool_size': 2})
+avg_pool = Operation.new('AveragePooling2D', {'pool_size': 2})
+global_pool = Operation.new('GlobalAveragePooling2D')
+
+
+class DebugSampler(Sampler):
+    def __init__(self):
+        self.iteration = 0
+
+    def choice(self, candidates, mutator, model, index):
+        idx = (self.iteration + index) % len(candidates)
+        return candidates[idx]
+
+    def mutation_start(self, mutator, model):
+        self.iteration += 1
+
+
+class DebugMutator(Mutator):
+    def mutate(self, model):
+        ops = [max_pool, avg_pool, global_pool]
+
+        pool1 = model.graphs['stem'].get_node_by_name('pool1')
+        pool1.update_operation(self.choice(ops))
+
+        pool2 = model.graphs['stem'].get_node_by_name('pool2')
+        pool2.update_operation(self.choice(ops))
+
+
+sampler = DebugSampler()
+mutator = DebugMutator()
+mutator.bind_sampler(sampler)
+
+
+json_path = Path(__file__).parent / 'mnist-tensorflow.json'
+ir = json.load(json_path.open())
+model0 = Model._load(ir)
+
+
+def test_dry_run():
+    candidates, _ = mutator.dry_run(model0)
+    assert len(candidates) == 2
+    assert candidates[0] == [max_pool, avg_pool, global_pool]
+    assert candidates[1] == [max_pool, avg_pool, global_pool]
+
+
+def test_mutation():
+    model1 = mutator.apply(model0)
+    assert _get_pools(model1) == (avg_pool, global_pool)
+
+    model2 = mutator.apply(model1)
+    assert _get_pools(model2) == (global_pool, max_pool)
+
+    assert model2.history == [model0, model1]
+    assert _get_pools(model0) == (max_pool, max_pool)
+    assert _get_pools(model1) == (avg_pool, global_pool)
+
+
+def _get_pools(model):
+    pool1 = model.graphs['stem'].get_node_by_name('pool1').operation
+    pool2 = model.graphs['stem'].get_node_by_name('pool2').operation
+    return pool1, pool2
+
+
+if __name__ == '__main__':
+    test_dry_run()
+    test_mutation()