Merge branch 'master' into nn-meter

examples/nas/multi-trial/nasbench101/base_ops.py

+import math
+
+import torch.nn as nn
+
+
+def truncated_normal_(tensor, mean=0, std=1):
+    # https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/15
+    size = tensor.shape
+    tmp = tensor.new_empty(size + (4,)).normal_()
+    valid = (tmp < 2) & (tmp > -2)
+    ind = valid.max(-1, keepdim=True)[1]
+    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
+    tensor.data.mul_(std).add_(mean)
+
+
+class ConvBnRelu(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0):
+        super(ConvBnRelu, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.conv_bn_relu = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                fan_in = m.kernel_size[0] * m.kernel_size[1] * m.in_channels
+                truncated_normal_(m.weight.data, mean=0., std=math.sqrt(1. / fan_in))
+            if isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        return self.conv_bn_relu(x)
+
+
+class Conv3x3BnRelu(ConvBnRelu):
+    def __init__(self, in_channels, out_channels):
+        super(Conv3x3BnRelu, self).__init__(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+
+class Conv1x1BnRelu(ConvBnRelu):
+    def __init__(self, in_channels, out_channels):
+        super(Conv1x1BnRelu, self).__init__(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+
+Projection = Conv1x1BnRelu

examples/nas/multi-trial/nasbench101/network.py

+import click
+import nni
+import nni.retiarii.evaluator.pytorch.lightning as pl
+import torch.nn as nn
+import torchmetrics
+from nni.retiarii import model_wrapper, serialize, serialize_cls
+from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
+from nni.retiarii.nn.pytorch import NasBench101Cell
+from nni.retiarii.strategy import Random
+from pytorch_lightning.callbacks import LearningRateMonitor
+from timm.optim import RMSpropTF
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torchvision import transforms
+from torchvision.datasets import CIFAR10
+
+from base_ops import Conv3x3BnRelu, Conv1x1BnRelu, Projection
+
+
+@model_wrapper
+class NasBench101(nn.Module):
+    def __init__(self,
+                 stem_out_channels: int = 128,
+                 num_stacks: int = 3,
+                 num_modules_per_stack: int = 3,
+                 max_num_vertices: int = 7,
+                 max_num_edges: int = 9,
+                 num_labels: int = 10,
+                 bn_eps: float = 1e-5,
+                 bn_momentum: float = 0.003):
+        super().__init__()
+
+        op_candidates = {
+            'conv3x3': lambda num_features: Conv3x3BnRelu(num_features, num_features),
+            'conv1x1': lambda num_features: Conv1x1BnRelu(num_features, num_features),
+            'maxpool': lambda num_features: nn.MaxPool2d(3, 1, 1)
+        }
+
+        # initial stem convolution
+        self.stem_conv = Conv3x3BnRelu(3, stem_out_channels)
+
+        layers = []
+        in_channels = out_channels = stem_out_channels
+        for stack_num in range(num_stacks):
+            if stack_num > 0:
+                downsample = nn.MaxPool2d(kernel_size=2, stride=2)
+                layers.append(downsample)
+                out_channels *= 2
+            for _ in range(num_modules_per_stack):
+                cell = NasBench101Cell(op_candidates, in_channels, out_channels,
+                                       lambda cin, cout: Projection(cin, cout),
+                                       max_num_vertices, max_num_edges, label='cell')
+                layers.append(cell)
+                in_channels = out_channels
+
+        self.features = nn.ModuleList(layers)
+        self.gap = nn.AdaptiveAvgPool2d(1)
+        self.classifier = nn.Linear(out_channels, num_labels)
+
+        for module in self.modules():
+            if isinstance(module, nn.BatchNorm2d):
+                module.eps = bn_eps
+                module.momentum = bn_momentum
+
+    def forward(self, x):
+        bs = x.size(0)
+        out = self.stem_conv(x)
+        for layer in self.features:
+            out = layer(out)
+        out = self.gap(out).view(bs, -1)
+        out = self.classifier(out)
+        return out
+
+    def reset_parameters(self):
+        for module in self.modules():
+            if isinstance(module, nn.BatchNorm2d):
+                module.eps = self.config.bn_eps
+                module.momentum = self.config.bn_momentum
+
+
+class AccuracyWithLogits(torchmetrics.Accuracy):
+    def update(self, pred, target):
+        return super().update(nn.functional.softmax(pred), target)
+
+
+@serialize_cls
+class NasBench101TrainingModule(pl.LightningModule):
+    def __init__(self, max_epochs=108, learning_rate=0.1, weight_decay=1e-4):
+        super().__init__()
+        self.save_hyperparameters('learning_rate', 'weight_decay', 'max_epochs')
+        self.criterion = nn.CrossEntropyLoss()
+        self.accuracy = AccuracyWithLogits()
+
+    def forward(self, x):
+        y_hat = self.model(x)
+        return y_hat
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        loss = self.criterion(y_hat, y)
+        self.log('train_loss', loss, prog_bar=True)
+        self.log('train_accuracy', self.accuracy(y_hat, y), prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        self.log('val_loss', self.criterion(y_hat, y), prog_bar=True)
+        self.log('val_accuracy', self.accuracy(y_hat, y), prog_bar=True)
+
+    def configure_optimizers(self):
+        optimizer = RMSpropTF(self.parameters(), lr=self.hparams.learning_rate,
+                              weight_decay=self.hparams.weight_decay,
+                              momentum=0.9, alpha=0.9, eps=1.0)
+        return {
+            'optimizer': optimizer,
+            'scheduler': CosineAnnealingLR(optimizer, self.hparams.max_epochs)
+        }
+
+    def on_validation_epoch_end(self):
+        nni.report_intermediate_result(self.trainer.callback_metrics['val_accuracy'].item())
+
+    def teardown(self, stage):
+        if stage == 'fit':
+            nni.report_final_result(self.trainer.callback_metrics['val_accuracy'].item())
+
+
+@click.command()
+@click.option('--epochs', default=108, help='Training length.')
+@click.option('--batch_size', default=256, help='Batch size.')
+@click.option('--port', default=8081, help='On which port the experiment is run.')
+def _multi_trial_test(epochs, batch_size, port):
+    # initalize dataset. Note that 50k+10k is used. It's a little different from paper
+    transf = [
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip()
+    ]
+    normalize = [
+        transforms.ToTensor(),
+        transforms.Normalize([0.49139968, 0.48215827, 0.44653124], [0.24703233, 0.24348505, 0.26158768])
+    ]
+    train_dataset = serialize(CIFAR10, 'data', train=True, download=True, transform=transforms.Compose(transf + normalize))
+    test_dataset = serialize(CIFAR10, 'data', train=False, transform=transforms.Compose(normalize))
+
+    # specify training hyper-parameters
+    training_module = NasBench101TrainingModule(max_epochs=epochs)
+    # FIXME: need to fix a bug in serializer for this to work
+    # lr_monitor = serialize(LearningRateMonitor, logging_interval='step')
+    trainer = pl.Trainer(max_epochs=epochs, gpus=1)
+    lightning = pl.Lightning(
+        lightning_module=training_module,
+        trainer=trainer,
+        train_dataloader=pl.DataLoader(train_dataset, batch_size=batch_size, shuffle=True),
+        val_dataloaders=pl.DataLoader(test_dataset, batch_size=batch_size),
+    )
+
+    strategy = Random()
+
+    model = NasBench101()
+
+    exp = RetiariiExperiment(model, lightning, [], strategy)
+
+    exp_config = RetiariiExeConfig('local')
+    exp_config.trial_concurrency = 2
+    exp_config.max_trial_number = 20
+    exp_config.trial_gpu_number = 1
+    exp_config.training_service.use_active_gpu = False
+
+    exp.run(exp_config, port)
+
+
+if __name__ == '__main__':
+    _multi_trial_test()

examples/nas/multi-trial/nasbench201/base_ops.py

+import torch
+import torch.nn as nn
+
+
+OPS_WITH_STRIDE = {
+    'none': lambda C_in, C_out, stride: Zero(C_in, C_out, stride),
+    'avg_pool_3x3': lambda C_in, C_out, stride: Pooling(C_in, C_out, stride, 'avg'),
+    'max_pool_3x3': lambda C_in, C_out, stride: Pooling(C_in, C_out, stride, 'max'),
+    'conv_3x3': lambda C_in, C_out, stride: ReLUConvBN(C_in, C_out, (3, 3), (stride, stride), (1, 1), (1, 1)),
+    'conv_1x1': lambda C_in, C_out, stride: ReLUConvBN(C_in, C_out, (1, 1), (stride, stride), (0, 0), (1, 1)),
+    'skip_connect': lambda C_in, C_out, stride: nn.Identity() if stride == 1 and C_in == C_out
+    else FactorizedReduce(C_in, C_out, stride),
+}
+
+PRIMITIVES = ['none', 'skip_connect', 'conv_1x1', 'conv_3x3', 'avg_pool_3x3']
+
+
+class ReLUConvBN(nn.Module):
+    def __init__(self, C_in, C_out, kernel_size, stride, padding, dilation):
+        super(ReLUConvBN, self).__init__()
+        self.op = nn.Sequential(
+            nn.ReLU(inplace=False),
+            nn.Conv2d(C_in, C_out, kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, bias=False),
+            nn.BatchNorm2d(C_out)
+        )
+
+    def forward(self, x):
+        return self.op(x)
+
+
+class SepConv(nn.Module):
+    def __init__(self, C_in, C_out, kernel_size, stride, padding, dilation):
+        super(SepConv, self).__init__()
+        self.op = nn.Sequential(
+            nn.ReLU(inplace=False),
+            nn.Conv2d(C_in, C_in, kernel_size=kernel_size, stride=stride,
+                      padding=padding, dilation=dilation, groups=C_in, bias=False),
+            nn.Conv2d(C_in, C_out, kernel_size=1, padding=0, bias=False),
+            nn.BatchNorm2d(C_out),
+        )
+
+    def forward(self, x):
+        return self.op(x)
+
+
+class Pooling(nn.Module):
+    def __init__(self, C_in, C_out, stride, mode):
+        super(Pooling, self).__init__()
+        if C_in == C_out:
+            self.preprocess = None
+        else:
+            self.preprocess = ReLUConvBN(C_in, C_out, 1, 1, 0, 1)
+        if mode == 'avg':
+            self.op = nn.AvgPool2d(3, stride=stride, padding=1, count_include_pad=False)
+        elif mode == 'max':
+            self.op = nn.MaxPool2d(3, stride=stride, padding=1)
+        else:
+            raise ValueError('Invalid mode={:} in Pooling'.format(mode))
+
+    def forward(self, x):
+        if self.preprocess:
+            x = self.preprocess(x)
+        return self.op(x)
+
+
+class Zero(nn.Module):
+    def __init__(self, C_in, C_out, stride):
+        super(Zero, self).__init__()
+        self.C_in = C_in
+        self.C_out = C_out
+        self.stride = stride
+        self.is_zero = True
+
+    def forward(self, x):
+        if self.C_in == self.C_out:
+            if self.stride == 1:
+                return x.mul(0.)
+            else:
+                return x[:, :, ::self.stride, ::self.stride].mul(0.)
+        else:
+            shape = list(x.shape)
+            shape[1] = self.C_out
+            zeros = x.new_zeros(shape, dtype=x.dtype, device=x.device)
+            return zeros
+
+
+class FactorizedReduce(nn.Module):
+    def __init__(self, C_in, C_out, stride):
+        super(FactorizedReduce, self).__init__()
+        self.stride = stride
+        self.C_in = C_in
+        self.C_out = C_out
+        self.relu = nn.ReLU(inplace=False)
+        if stride == 2:
+            C_outs = [C_out // 2, C_out - C_out // 2]
+            self.convs = nn.ModuleList()
+            for i in range(2):
+                self.convs.append(nn.Conv2d(C_in, C_outs[i], 1, stride=stride, padding=0, bias=False))
+            self.pad = nn.ConstantPad2d((0, 1, 0, 1), 0)
+        else:
+            raise ValueError('Invalid stride : {:}'.format(stride))
+        self.bn = nn.BatchNorm2d(C_out)
+
+    def forward(self, x):
+        x = self.relu(x)
+        y = self.pad(x)
+        out = torch.cat([self.convs[0](x), self.convs[1](y[:, :, 1:, 1:])], dim=1)
+        out = self.bn(out)
+        return out
+
+
+class ResNetBasicblock(nn.Module):
+    def __init__(self, inplanes, planes, stride):
+        super(ResNetBasicblock, self).__init__()
+        assert stride == 1 or stride == 2, 'invalid stride {:}'.format(stride)
+        self.conv_a = ReLUConvBN(inplanes, planes, 3, stride, 1, 1)
+        self.conv_b = ReLUConvBN(planes, planes, 3, 1, 1, 1)
+        if stride == 2:
+            self.downsample = nn.Sequential(
+                nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
+                nn.Conv2d(inplanes, planes, kernel_size=1, stride=1, padding=0, bias=False))
+        elif inplanes != planes:
+            self.downsample = ReLUConvBN(inplanes, planes, 1, 1, 0, 1)
+        else:
+            self.downsample = None
+        self.in_dim = inplanes
+        self.out_dim = planes
+        self.stride = stride
+        self.num_conv = 2
+
+    def forward(self, inputs):
+        basicblock = self.conv_a(inputs)
+        basicblock = self.conv_b(basicblock)
+
+        if self.downsample is not None:
+            inputs = self.downsample(inputs)  # residual
+        return inputs + basicblock

examples/nas/multi-trial/nasbench201/network.py

+import click
+import nni
+import nni.retiarii.evaluator.pytorch.lightning as pl
+import torch.nn as nn
+import torchmetrics
+from nni.retiarii import model_wrapper, serialize, serialize_cls
+from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
+from nni.retiarii.nn.pytorch import NasBench201Cell
+from nni.retiarii.strategy import Random
+from pytorch_lightning.callbacks import LearningRateMonitor
+from timm.optim import RMSpropTF
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torchvision import transforms
+from torchvision.datasets import CIFAR100
+
+from base_ops import ResNetBasicblock, PRIMITIVES, OPS_WITH_STRIDE
+
+
+@model_wrapper
+class NasBench201(nn.Module):
+    def __init__(self,
+                 stem_out_channels: int = 16,
+                 num_modules_per_stack: int = 5,
+                 num_labels: int = 100):
+        super().__init__()
+        self.channels = C = stem_out_channels
+        self.num_modules = N = num_modules_per_stack
+        self.num_labels = num_labels
+
+        self.stem = nn.Sequential(
+            nn.Conv2d(3, C, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(C)
+        )
+
+        layer_channels = [C] * N + [C * 2] + [C * 2] * N + [C * 4] + [C * 4] * N
+        layer_reductions = [False] * N + [True] + [False] * N + [True] + [False] * N
+
+        C_prev = C
+        self.cells = nn.ModuleList()
+        for C_curr, reduction in zip(layer_channels, layer_reductions):
+            if reduction:
+                cell = ResNetBasicblock(C_prev, C_curr, 2)
+            else:
+                cell = NasBench201Cell({prim: lambda C_in, C_out: OPS_WITH_STRIDE[prim](C_in, C_out, 1) for prim in PRIMITIVES},
+                                       C_prev, C_curr, label='cell')
+            self.cells.append(cell)
+            C_prev = C_curr
+
+        self.lastact = nn.Sequential(
+            nn.BatchNorm2d(C_prev),
+            nn.ReLU(inplace=True)
+        )
+        self.global_pooling = nn.AdaptiveAvgPool2d(1)
+        self.classifier = nn.Linear(C_prev, self.num_labels)
+
+    def forward(self, inputs):
+        feature = self.stem(inputs)
+        for cell in self.cells:
+            feature = cell(feature)
+
+        out = self.lastact(feature)
+        out = self.global_pooling(out)
+        out = out.view(out.size(0), -1)
+        logits = self.classifier(out)
+
+        return logits
+
+
+class AccuracyWithLogits(torchmetrics.Accuracy):
+    def update(self, pred, target):
+        return super().update(nn.functional.softmax(pred), target)
+
+
+@serialize_cls
+class NasBench201TrainingModule(pl.LightningModule):
+    def __init__(self, max_epochs=200, learning_rate=0.1, weight_decay=5e-4):
+        super().__init__()
+        self.save_hyperparameters('learning_rate', 'weight_decay', 'max_epochs')
+        self.criterion = nn.CrossEntropyLoss()
+        self.accuracy = AccuracyWithLogits()
+
+    def forward(self, x):
+        y_hat = self.model(x)
+        return y_hat
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        loss = self.criterion(y_hat, y)
+        self.log('train_loss', loss, prog_bar=True)
+        self.log('train_accuracy', self.accuracy(y_hat, y), prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        self.log('val_loss', self.criterion(y_hat, y), prog_bar=True)
+        self.log('val_accuracy', self.accuracy(y_hat, y), prog_bar=True)
+
+    def configure_optimizers(self):
+        optimizer = RMSpropTF(self.parameters(), lr=self.hparams.learning_rate,
+                              weight_decay=self.hparams.weight_decay,
+                              momentum=0.9, alpha=0.9, eps=1.0)
+        return {
+            'optimizer': optimizer,
+            'scheduler': CosineAnnealingLR(optimizer, self.hparams.max_epochs)
+        }
+
+    def on_validation_epoch_end(self):
+        nni.report_intermediate_result(self.trainer.callback_metrics['val_accuracy'].item())
+
+    def teardown(self, stage):
+        if stage == 'fit':
+            nni.report_final_result(self.trainer.callback_metrics['val_accuracy'].item())
+
+
+@click.command()
+@click.option('--epochs', default=12, help='Training length.')
+@click.option('--batch_size', default=256, help='Batch size.')
+@click.option('--port', default=8081, help='On which port the experiment is run.')
+def _multi_trial_test(epochs, batch_size, port):
+    # initalize dataset. Note that 50k+10k is used. It's a little different from paper
+    transf = [
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip()
+    ]
+    normalize = [
+        transforms.ToTensor(),
+        transforms.Normalize([x / 255 for x in [129.3, 124.1, 112.4]], [x / 255 for x in [68.2, 65.4, 70.4]])
+    ]
+    train_dataset = serialize(CIFAR100, 'data', train=True, download=True, transform=transforms.Compose(transf + normalize))
+    test_dataset = serialize(CIFAR100, 'data', train=False, transform=transforms.Compose(normalize))
+
+    # specify training hyper-parameters
+    training_module = NasBench201TrainingModule(max_epochs=epochs)
+    # FIXME: need to fix a bug in serializer for this to work
+    # lr_monitor = serialize(LearningRateMonitor, logging_interval='step')
+    trainer = pl.Trainer(max_epochs=epochs, gpus=1)
+    lightning = pl.Lightning(
+        lightning_module=training_module,
+        trainer=trainer,
+        train_dataloader=pl.DataLoader(train_dataset, batch_size=batch_size, shuffle=True),
+        val_dataloaders=pl.DataLoader(test_dataset, batch_size=batch_size),
+    )
+
+    strategy = Random()
+
+    model = NasBench201()
+
+    exp = RetiariiExperiment(model, lightning, [], strategy)
+
+    exp_config = RetiariiExeConfig('local')
+    exp_config.trial_concurrency = 2
+    exp_config.max_trial_number = 20
+    exp_config.trial_gpu_number = 1
+    exp_config.training_service.use_active_gpu = False
+
+    exp.run(exp_config, port)
+
+
+if __name__ == '__main__':
+    _multi_trial_test()

examples/trials/mnist-pytorch/config_detailed.yml

@@ -30,7 +30,7 @@ trialConcurrency: 4             # Run 4 trials concurrently.
 maxTrialNumber: 10              # Generate at most 10 trials.
 maxExperimentDuration: 1h       # Stop generating trials after 1 hour.
 
-tuner:                          # Configure the tuning alogrithm.
+tuner:                          # Configure the tuning algorithm.
   name: TPE                     # Supported algorithms: TPE, Random, Anneal, Evolution, GridSearch, GPTuner, PBTTuner, etc.
                                 #   Full list:  https://nni.readthedocs.io/en/latest/Tuner/BuiltinTuner.html
   classArgs:                    # Algorithm specific arguments. See the tuner's doc for details.

examples/trials/mnist-tfv2/config_detailed.yml

@@ -33,7 +33,7 @@ trialConcurrency: 4             # Run 4 trials concurrently.
 maxTrialNumber: 10              # Generate at most 10 trials.
 maxExperimentDuration: 1h       # Stop generating trials after 1 hour.
 
-tuner:                          # Configure the tuning alogrithm.
+tuner:                          # Configure the tuning algorithm.
   name: TPE                     # Supported algorithms: TPE, Random, Anneal, Evolution, GridSearch, GPTuner, PBTTuner, etc.
                                 #   Full list:  https://nni.readthedocs.io/en/latest/Tuner/BuiltinTuner.html
   classArgs:                    # Algorithm specific arguments. See the tuner's doc for details.

examples/trials/pix2pix-pytorch/.gitignore

+# datasets
+data/
+
+# pix2pix library
+pix2pixlib/
\ No newline at end of file

examples/trials/pix2pix-pytorch/base_params.py

+
+
+def get_base_params(dataset_name, checkpoint_dir):
+    params = {}
+
+    # change name and gpuid later
+    basic_params = {'dataset': dataset_name,
+                    'dataroot': './data/' + dataset_name,
+                    'name': '',
+                    'gpu_ids': [0],  
+                    'checkpoints_dir': checkpoint_dir,
+                    'verbose': False,
+                    'print_freq': 100
+                    }
+    params.update(basic_params)
+    
+    dataset_params = {'dataset_mode': 'aligned',
+                      'direction': 'BtoA',
+                      'num_threads': 4,
+                      'max_dataset_size': float('inf'),
+                      'preprocess': 'resize_and_crop',
+                      'display_winsize': 256,
+                      'input_nc': 3,
+                      'output_nc': 3}
+    params.update(dataset_params)  
+
+    model_params = {'model': 'pix2pix',
+                    # 'ngf': 64,
+                    # 'ndf': 64,
+                    # 'netD': 'basic',
+                    # 'netG': 'unet_256',
+                    'n_layers_D': 3, 
+                    # 'norm': 'batch',
+                    # 'gan_mode': 'lsgan',
+                    # 'init_type': 'normal',
+                    'init_gain': 0.02,
+                    'no_dropout': False}
+    params.update(model_params)
+    
+    train_params = {'phase': 'train',
+                    'isTrain': True,
+                    'serial_batches': False,
+                    'load_size': 286,
+                    'crop_size': 256,
+                    'no_flip': False,
+                    # 'batch_size': 1,
+                    # 'beta1': 0.5,
+                    'pool_size': 0,
+                    # 'lr_policy': 'linear',
+                    'lr_decay_iters': 50,
+                    #'lr': 0.0002,
+                    # 'lambda_L1': 100,
+                    'epoch_count': 1,
+                    # 'n_epochs': 10,           # 100
+                    # 'n_epochs_decay': 0,      # 100
+                    'continue_train': False}
+    train_params.update(params)
+    
+    test_params = {'phase': 'test',
+                   'isTrain': False,
+                   'load_iter': -1,
+                   'epoch': 'latest',
+                   'load_size': 256,
+                   'crop_size': 256,
+                   # 'batch_size': 1,
+                   'serial_batches': True,
+                   'no_flip': True,
+                   'eval': True}
+    test_params.update(params)
+    
+    return train_params, test_params
+

examples/trials/pix2pix-pytorch/config.yml

+experimentName: example_pix2pix 
+searchSpace:
+  ngf:
+    _type: choice
+    _value: [16, 32, 64, 128]
+  ndf:
+    _type: choice
+    _value: [16, 32, 64, 128]
+  netG:
+    _type: choice
+    _value: ["unet_256", "resnet_9blocks"]
+  netD:
+    _type: choice
+    _value: ["basic", "pixel", "n_layers"]
+  norm:
+    _type: choice
+    _value: ["batch", "instance", "none"]
+  init_type:
+    _type: choice
+    _value: ["xavier", "normal", "kaiming", "orthogonal"]
+  lr:
+    _type: choice
+    _value: [0.0001, 0.0002, 0.0005, 0.001, 0.005, 0.01, 0.1]
+  beta1:
+    _type: uniform
+    _value: [0, 1]
+  lr_policy:
+    _type: choice
+    _value: ["linear", "step", "plateau", "cosine"]
+  gan_mode:
+    _type: choice
+    _value: ["vanilla", "lsgan", "wgangp"]
+  lambda_L1:
+    _type: choice
+    _value: [1, 5, 10, 100, 250, 500]
+trainingService:
+  platform: local
+  useActiveGpu: true
+  gpuIndices: '0'
+trialCodeDirectory: .
+trialCommand: python3 pix2pix.py
+trialConcurrency: 1
+trialGpuNumber: 1
+tuner:
+  name: TPE
+  classArgs:
+    optimize_mode: minimize
\ No newline at end of file

examples/trials/pix2pix-pytorch/pix2pix.py

+import sys
+sys.path.insert(0, './pix2pixlib')
+
+import os
+import pathlib
+import logging
+import time
+import argparse
+from collections import namedtuple
+import numpy as np
+import torch
+import torch.utils.data as data
+import nni
+from nni.utils import merge_parameter
+from pix2pixlib.data.aligned_dataset import AlignedDataset
+from pix2pixlib.models.pix2pix_model import Pix2PixModel
+from base_params import get_base_params
+
+
+_logger = logging.getLogger('example_pix2pix')
+
+
+class CustomDatasetDataLoader():
+    """Wrapper class of Dataset class that performs multi-threaded data loading"""
+
+    def __init__(self, opt, ds):
+        """Initialize this class
+        Step 1: create a dataset instance given the name [dataset_mode]
+        Step 2: create a multi-threaded data loader.
+        """
+        self.opt = opt
+        self.dataset = ds
+        self.dataloader = data.DataLoader(self.dataset,
+                                          batch_size=opt.batch_size,
+                                          shuffle=not opt.serial_batches,
+                                          num_workers=int(opt.num_threads))
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        """Return the number of data in the dataset"""
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        """Return a batch of data"""
+        for i, data in enumerate(self.dataloader):
+            if i * self.opt.batch_size >= self.opt.max_dataset_size:
+                break
+            yield data
+
+
+def download_dataset(dataset_name):
+    # code adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
+    assert(dataset_name in ['facades', 'night2day', 'edges2handbags', 'edges2shoes', 'maps'])
+    if os.path.exists('./data/' + dataset_name):
+        _logger.info("Already downloaded dataset " + dataset_name)
+    else:
+        _logger.info("Downloading dataset " + dataset_name)
+        if not os.path.exists('./data/'):
+            pathlib.Path('./data/').mkdir(parents=True, exist_ok=True)
+        pathlib.Path('./data/' + dataset_name).mkdir(parents=True, exist_ok=True)
+        URL = 'http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/{}.tar.gz'.format(dataset_name)
+        TAR_FILE = './data/{}.tar.gz'.format(dataset_name)
+        TARGET_DIR = './data/{}/'.format(dataset_name)
+        os.system('wget -N {} -O {}'.format(URL, TAR_FILE))
+        pathlib.Path(TARGET_DIR).mkdir(parents=True, exist_ok=True)
+        os.system('tar -zxvf {} -C ./data/'.format(TAR_FILE))
+        os.system('rm ' + TAR_FILE)        
+    
+
+def setup_trial_checkpoint_dir():
+    checkpoint_dir = os.environ['NNI_OUTPUT_DIR'] + '/checkpoints/'
+    pathlib.Path(checkpoint_dir).mkdir(parents=True, exist_ok=True)
+    return checkpoint_dir
+        
+
+def parse_args():
+    # Settings that may be overrided by parameters from nni
+    parser = argparse.ArgumentParser(description='PyTorch Pix2pix Example')
+    parser.add_argument('--ngf', type=int, default=64,
+                        help='# of generator filters in the last conv layer')
+    parser.add_argument('--ndf', type=int, default=64,
+                        help='# of discriminator filters in the first conv layer')
+    parser.add_argument('--netD', type=str, default='basic',
+                        help='specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator')
+    parser.add_argument('--netG', type=str, default='resnet_9blocks',
+                        help='specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]')
+    parser.add_argument('--init_type', type=str, default='normal',
+                        help='network initialization [normal | xavier | kaiming | orthogonal]')
+    parser.add_argument('--beta1', type=float, default=0.5,
+                        help='momentum term of adam')
+    parser.add_argument('--lr', type=float, default=0.0002,
+                        help='initial learning rate for adam')
+    parser.add_argument('--lr_policy', type=str, default='linear',
+                        help='learning rate policy. [linear | step | plateau | cosine]')
+    parser.add_argument('--gan_mode', type=str, default='lsgan',
+                        help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.')
+    parser.add_argument('--norm', type=str, default='instance',
+                        help='instance normalization or batch normalization [instance | batch | none]')
+    parser.add_argument('--lambda_L1', type=float, default=100,
+                        help='weight of L1 loss in the generator objective')
+    
+    # Additional training settings 
+    parser.add_argument('--batch_size', type=int, default=1,
+                        help='input batch size for training (default: 1)')
+    parser.add_argument('--n_epochs', type=int, default=100,
+                        help='number of epochs with the initial learning rate')
+    parser.add_argument('--n_epochs_decay', type=int, default=100,
+                        help='number of epochs to linearly decay learning rate to zero')
+    
+    args, _ = parser.parse_known_args()
+    return args
+  
+
+def evaluate_L1(config, model, dataset):
+    if config.eval:
+        model.eval()
+    scores = []
+    for i, data in enumerate(dataset):
+        model.set_input(data)  # unpack data from data loader
+        model.test()           # run inference
+        visuals = model.get_current_visuals()
+        score = torch.mean(torch.abs(visuals['fake_B']-visuals['real_B'])).detach().cpu().numpy()
+        scores.append(score)
+    return np.mean(np.array(scores))
+            
+
+def main(dataset_name, train_params, test_params):
+    download_dataset(dataset_name)
+    
+    train_config = namedtuple('Struct', train_params.keys())(*train_params.values())
+    test_config = namedtuple('Struct', test_params.keys())(*test_params.values())
+    
+    train_dataset, test_dataset = AlignedDataset(train_config), AlignedDataset(test_config)
+    print(train_dataset, train_config)
+    train_dataset = CustomDatasetDataLoader(train_config, train_dataset)
+    test_dataset = CustomDatasetDataLoader(test_config, test_dataset)
+    _logger.info('Number of training images = {}'.format(len(train_dataset)))
+    _logger.info('Number of testing images = {}'.format(len(test_dataset)))    
+
+    model = Pix2PixModel(train_config)
+    model.setup(train_config)
+
+    # training 
+    total_iters = 0                # the total number of training iterations
+    for epoch in range(train_config.epoch_count, train_config.n_epochs + train_config.n_epochs_decay + 1):
+        _logger.info('Training epoch {}'.format(epoch))
+        epoch_start_time = time.time()  # timer for entire epoch
+        iter_data_time = time.time()    # timer for data loading per iteration
+        epoch_iter = 0
+        model.update_learning_rate()
+        for i, data in enumerate(train_dataset):  # inner loop within one epoch
+            iter_start_time = time.time()  # timer for computation per iteration
+            if total_iters % train_config.print_freq == 0:
+                t_data = iter_start_time - iter_data_time
+            total_iters += train_config.batch_size
+            epoch_iter += train_config.batch_size
+            model.set_input(data)         # unpack data from dataset and apply preprocessing
+            model.optimize_parameters()   # calculate loss functions, get gradients, update network weights
+            iter_data_time = time.time()
+        _logger.info('End of epoch {} / {} \t Time Taken: {} sec'.format(epoch, train_config.n_epochs + train_config.n_epochs_decay, time.time() - epoch_start_time))
+    
+    model.save_networks('latest')
+    _logger.info("Training done. Saving the final model.")
+    
+    l1_score = evaluate_L1(test_config, model, test_dataset)
+    _logger.info("The final L1 loss the test set is {}".format(l1_score))  
+    nni.report_final_result(l1_score)
+
+    
+if __name__ == '__main__':
+    dataset_name = 'facades'
+
+    checkpoint_dir = setup_trial_checkpoint_dir()
+
+    params_from_cl = vars(parse_args())
+    params_for_tuning = nni.get_next_parameter()       
+    train_params, test_params = get_base_params(dataset_name, checkpoint_dir)
+    train_params.update(params_from_cl)
+    test_params.update(params_from_cl)
+    train_params = merge_parameter(train_params, params_for_tuning)
+        
+    main(dataset_name, train_params, test_params)
+    

examples/trials/pix2pix-pytorch/setup.sh

+#!/bin/bash
+
+# download pix2pix repository
+if [ ! -d './pix2pixlib' ] ; then
+    git clone https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix.git pix2pixlib
+fi

examples/trials/pix2pix-pytorch/test.py

+import sys
+sys.path.insert(0, './pix2pixlib')
+
+import os
+import pathlib
+import logging
+import argparse
+import json
+from collections import namedtuple
+from PIL import Image
+import numpy as np
+import torch
+from nni.utils import merge_parameter
+from pix2pixlib.data.aligned_dataset import AlignedDataset
+from pix2pixlib.data import CustomDatasetDataLoader
+from pix2pixlib.models.pix2pix_model import Pix2PixModel
+from pix2pixlib.util.util import tensor2im
+from base_params import get_base_params
+
+
+_logger = logging.getLogger('example_pix2pix')
+
+
+def download_dataset(dataset_name):
+    # code adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
+    assert(dataset_name in ['facades', 'night2day', 'edges2handbags', 'edges2shoes', 'maps'])
+    if os.path.exists('./data/' + dataset_name):
+        _logger.info("Already downloaded dataset " + dataset_name)
+    else:
+        _logger.info("Downloading dataset " + dataset_name)
+        if not os.path.exists('./data/'):
+            pathlib.Path('./data/').mkdir(parents=True, exist_ok=True)
+        pathlib.Path('./data/' + dataset_name).mkdir(parents=True, exist_ok=True)
+        URL = 'http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/{}.tar.gz'.format(dataset_name)
+        TAR_FILE = './data/{}.tar.gz'.format(dataset_name)
+        TARGET_DIR = './data/{}/'.format(dataset_name)
+        os.system('wget -N {} -O {}'.format(URL, TAR_FILE))
+        pathlib.Path(TARGET_DIR).mkdir(parents=True, exist_ok=True)
+        os.system('tar -zxvf {} -C ./data/'.format(TAR_FILE))
+        os.system('rm ' + TAR_FILE)        
+
+        
+def parse_args():
+    parser = argparse.ArgumentParser(description='PyTorch Pix2pix Example')
+
+    # required arguments
+    parser.add_argument('-c', '--checkpoint', type=str, required=True,
+                        help='Checkpoint directory')
+    parser.add_argument('-p', '--parameter_cfg', type=str, required=True,
+                        help='parameter.cfg file generated by nni trial')
+    parser.add_argument('-d', '--dataset', type=str, required=True,
+                        help='dataset name (facades, night2day, edges2handbags, edges2shoes, maps)')
+    parser.add_argument('-o', '--output_dir', type=str, required=True,
+                        help='Where to save the test results')
+    
+    # Settings that may be overrided by parameters from nni
+    parser.add_argument('--ngf', type=int, default=64, 
+                        help='# of generator filters in the last conv layer')
+    parser.add_argument('--ndf', type=int, default=64,
+                        help='# of discriminator filters in the first conv layer')
+    parser.add_argument('--netD', type=str, default='basic',
+                        help='specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator')
+    parser.add_argument('--netG', type=str, default='resnet_9blocks',
+                        help='specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]')
+    parser.add_argument('--init_type', type=str, default='normal',
+                        help='network initialization [normal | xavier | kaiming | orthogonal]')
+    parser.add_argument('--beta1', type=float, default=0.5,
+                        help='momentum term of adam')
+    parser.add_argument('--lr', type=float, default=0.0002,
+                        help='initial learning rate for adam')
+    parser.add_argument('--lr_policy', type=str, default='linear',
+                        help='learning rate policy. [linear | step | plateau | cosine]')
+    parser.add_argument('--gan_mode', type=str, default='lsgan',
+                        help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.')
+    parser.add_argument('--norm', type=str, default='instance',
+                        help='instance normalization or batch normalization [instance | batch | none]')
+    parser.add_argument('--lambda_L1', type=float, default=100,
+                        help='weight of L1 loss in the generator objective')
+    
+    # Additional training settings 
+    parser.add_argument('--batch_size', type=int, default=1,
+                        help='input batch size for training (default: 1)')
+    parser.add_argument('--n_epochs', type=int, default=100,
+                        help='number of epochs with the initial learning rate')
+    parser.add_argument('--n_epochs_decay', type=int, default=100,
+                        help='number of epochs to linearly decay learning rate to zero')
+    
+    args, _ = parser.parse_known_args()
+    return args
+
+
+def main(test_params):
+    test_config = namedtuple('Struct', test_params.keys())(*test_params.values())
+    assert os.path.exists(test_config.checkpoint), "Checkpoint does not exist"
+
+    download_dataset(test_config.dataset)
+    
+    test_dataset = AlignedDataset(test_config)
+    test_dataset = CustomDatasetDataLoader(test_config, test_dataset)
+    _logger.info('Number of testing images = {}'.format(len(test_dataset)))    
+
+    model = Pix2PixModel(test_config)
+    model.setup(test_config)
+
+    if test_config.eval:
+        model.eval()
+
+    for i, data in enumerate(test_dataset):
+        print('Testing on {} image {}'.format(test_config.dataset, i), end='\r')
+        model.set_input(data)  
+        model.test()
+
+        visuals = model.get_current_visuals()
+        cur_input = tensor2im(visuals['real_A'])
+        cur_label = tensor2im(visuals['real_B'])
+        cur_output = tensor2im(visuals['fake_B'])
+        
+        image_name = '{}_test_{}.png'.format(test_config.dataset, i)
+        Image.fromarray(cur_input).save(os.path.join(test_config.output_dir, 'input', image_name))
+        Image.fromarray(cur_label).save(os.path.join(test_config.output_dir, 'label', image_name))
+        Image.fromarray(cur_output).save(os.path.join(test_config.output_dir, 'output', image_name))
+
+    _logger.info("Images successfully saved to " + test_config.output_dir)
+
+    
+if __name__ == '__main__':
+    params_from_cl = vars(parse_args())
+    _, test_params = get_base_params(params_from_cl['dataset'], params_from_cl['checkpoint'])  
+    test_params.update(params_from_cl)
+
+    with open(test_params['parameter_cfg'], 'r') as f:
+        params_from_nni = json.loads(f.readline().strip())['parameters']
+    test_params = merge_parameter(test_params, params_from_nni)
+
+    pathlib.Path(params_from_cl['output_dir'] + '/input').mkdir(parents=True, exist_ok=True)
+    pathlib.Path(params_from_cl['output_dir'] + '/label').mkdir(parents=True, exist_ok=True)
+    pathlib.Path(params_from_cl['output_dir'] + '/output').mkdir(parents=True, exist_ok=True)
+
+    main(test_params)
+    

nni/algorithms/compression/pytorch/pruning/auto_compress_pruner.py

@@ -11,7 +11,7 @@ from nni.utils import OptimizeMode
 from nni.compression.pytorch import ModelSpeedup
 
 from nni.compression.pytorch.compressor import Pruner
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from .simulated_annealing_pruner import SimulatedAnnealingPruner
 from .iterative_pruner import ADMMPruner
 
@@ -130,16 +130,18 @@ class AutoCompressPruner(Pruner):
         """
 
         if self._base_algo == 'level':
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 Optional('op_types'): [str],
                 Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
         elif self._base_algo in ['l1', 'l2', 'fpgm']:
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 'op_types': ['Conv2d'],
-                Optional('op_names'): [str]
+                Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
 
         schema.validate(config_list)

nni/algorithms/compression/pytorch/pruning/dependency_aware_pruner.py

@@ -2,10 +2,10 @@
 # Licensed under the MIT license.
 
 import logging
-from schema import And, Optional, SchemaError
+from schema import And, Optional
 from nni.common.graph_utils import TorchModuleGraph
 from nni.compression.pytorch.utils.shape_dependency import ChannelDependency, GroupDependency
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from nni.compression.pytorch.compressor import Pruner
 from .constants import MASKER_DICT
 
@@ -82,7 +82,7 @@ class DependencyAwarePruner(Pruner):
             self._dependency_update_mask()
 
     def validate_config(self, model, config_list):
-        schema = CompressorSchema([{
+        schema = PrunerSchema([{
             Optional('sparsity'): And(float, lambda n: 0 < n < 1),
             Optional('op_types'): ['Conv2d'],
             Optional('op_names'): [str],
@@ -90,9 +90,6 @@ class DependencyAwarePruner(Pruner):
         }], model, logger)
 
         schema.validate(config_list)
-        for config in config_list:
-            if 'exclude' not in config and 'sparsity' not in config:
-                raise SchemaError('Either sparisty or exclude must be specified!')
 
     def _supported_dependency_aware(self):
         raise NotImplementedError

nni/algorithms/compression/pytorch/pruning/iterative_pruner.py

@@ -5,7 +5,7 @@ import logging
 import copy
 import torch
 from schema import And, Optional
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from .constants import MASKER_DICT
 from .dependency_aware_pruner import DependencyAwarePruner
 
@@ -138,10 +138,11 @@ class AGPPruner(IterativePruner):
         config_list : list
             List on pruning configs
         """
-        schema = CompressorSchema([{
-            'sparsity': And(float, lambda n: 0 <= n <= 1),
+        schema = PrunerSchema([{
+            Optional('sparsity'): And(float, lambda n: 0 <= n <= 1),
             Optional('op_types'): [str],
-            Optional('op_names'): [str]
+            Optional('op_names'): [str],
+            Optional('exclude'): bool
         }], model, logger)
 
         schema.validate(config_list)
@@ -300,16 +301,18 @@ class ADMMPruner(IterativePruner):
         """
 
         if self._base_algo == 'level':
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 Optional('op_types'): [str],
                 Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, logger)
         elif self._base_algo in ['l1', 'l2', 'fpgm']:
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 'op_types': ['Conv2d'],
-                Optional('op_names'): [str]
+                Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, logger)
 
         schema.validate(config_list)
@@ -436,10 +439,11 @@ class SlimPruner(IterativePruner):
         self.patch_optimizer_before(self._callback)
 
     def validate_config(self, model, config_list):
-        schema = CompressorSchema([{
-            'sparsity': And(float, lambda n: 0 < n < 1),
+        schema = PrunerSchema([{
+            Optional('sparsity'): And(float, lambda n: 0 < n < 1),
             'op_types': ['BatchNorm2d'],
-            Optional('op_names'): [str]
+            Optional('op_names'): [str],
+            Optional('exclude'): bool
         }], model, logger)
 
         schema.validate(config_list)

nni/algorithms/compression/pytorch/pruning/lottery_ticket.py

@@ -5,7 +5,7 @@ import copy
 import logging
 import torch
 from schema import And, Optional
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from nni.compression.pytorch.compressor import Pruner
 from .finegrained_pruning_masker import LevelPrunerMasker
 
@@ -56,11 +56,12 @@ class LotteryTicketPruner(Pruner):
                 - prune_iterations : The number of rounds for the iterative pruning.
                 - sparsity : The final sparsity when the compression is done.
         """
-        schema = CompressorSchema([{
-            'sparsity': And(float, lambda n: 0 < n < 1),
+        schema = PrunerSchema([{
+            Optional('sparsity'): And(float, lambda n: 0 < n < 1),
             'prune_iterations': And(int, lambda n: n > 0),
             Optional('op_types'): [str],
-            Optional('op_names'): [str]
+            Optional('op_names'): [str],
+            Optional('exclude'): bool
         }], model, logger)
 
         schema.validate(config_list)

nni/algorithms/compression/pytorch/pruning/net_adapt_pruner.py

@@ -11,7 +11,7 @@ from schema import And, Optional
 from nni.utils import OptimizeMode
 
 from nni.compression.pytorch.compressor import Pruner
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from nni.compression.pytorch.utils.num_param_counter import get_total_num_weights
 from .constants_pruner import PRUNER_DICT
 
@@ -120,16 +120,18 @@ class NetAdaptPruner(Pruner):
         """
 
         if self._base_algo == 'level':
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 Optional('op_types'): [str],
                 Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
         elif self._base_algo in ['l1', 'l2', 'fpgm']:
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 'op_types': ['Conv2d'],
-                Optional('op_names'): [str]
+                Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
 
         schema.validate(config_list)

nni/algorithms/compression/pytorch/pruning/one_shot_pruner.py

@@ -4,7 +4,7 @@
 import logging
 from schema import And, Optional
 
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from .dependency_aware_pruner import DependencyAwarePruner
 
 __all__ = ['LevelPruner', 'L1FilterPruner', 'L2FilterPruner', 'FPGMPruner']
@@ -48,10 +48,11 @@ class OneshotPruner(DependencyAwarePruner):
         config_list : list
             List on pruning configs
         """
-        schema = CompressorSchema([{
-            'sparsity': And(float, lambda n: 0 < n < 1),
+        schema = PrunerSchema([{
+            Optional('sparsity'): And(float, lambda n: 0 < n < 1),
             Optional('op_types'): [str],
-            Optional('op_names'): [str]
+            Optional('op_names'): [str],
+            Optional('exclude'): bool
         }], model, logger)
 
         schema.validate(config_list)

nni/algorithms/compression/pytorch/pruning/sensitivity_pruner.py

@@ -9,7 +9,7 @@ import torch
 
 from schema import And, Optional
 from nni.compression.pytorch.compressor import Pruner
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from nni.compression.pytorch.utils.sensitivity_analysis import SensitivityAnalysis
 
 from .constants_pruner import PRUNER_DICT
@@ -146,16 +146,18 @@ class SensitivityPruner(Pruner):
         """
 
         if self.base_algo == 'level':
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 Optional('op_types'): [str],
                 Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
         elif self.base_algo in ['l1', 'l2', 'fpgm']:
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 'op_types': ['Conv2d'],
-                Optional('op_names'): [str]
+                Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
 
         schema.validate(config_list)

nni/algorithms/compression/pytorch/pruning/simulated_annealing_pruner.py

@@ -13,7 +13,7 @@ from schema import And, Optional
 from nni.utils import OptimizeMode
 
 from nni.compression.pytorch.compressor import Pruner
-from nni.compression.pytorch.utils.config_validation import CompressorSchema
+from nni.compression.pytorch.utils.config_validation import PrunerSchema
 from .constants_pruner import PRUNER_DICT
 
 
@@ -115,16 +115,18 @@ class SimulatedAnnealingPruner(Pruner):
         """
 
         if self._base_algo == 'level':
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 Optional('op_types'): [str],
                 Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
         elif self._base_algo in ['l1', 'l2', 'fpgm']:
-            schema = CompressorSchema([{
-                'sparsity': And(float, lambda n: 0 < n < 1),
+            schema = PrunerSchema([{
+                Optional('sparsity'): And(float, lambda n: 0 < n < 1),
                 'op_types': ['Conv2d'],
-                Optional('op_names'): [str]
+                Optional('op_names'): [str],
+                Optional('exclude'): bool
             }], model, _logger)
 
         schema.validate(config_list)