Merge pull request #239 from microsoft/master

merge master

examples/nas/textnas/retrain.py

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import sys
+import os
+import logging
+import pickle
+import shutil
+import random
+import math
+
+import time
+import datetime
+import argparse
+import distutils.util
+
+import numpy as np
+import torch
+from torch import nn
+from torch import optim
+from torch.utils.data import DataLoader
+import torch.nn.functional as Func
+
+from model import Model
+from nni.nas.pytorch.fixed import apply_fixed_architecture
+from dataloader import read_data_sst
+
+
+logger = logging.getLogger("nni.textnas")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--reset_output_dir",
+        type=distutils.util.strtobool,
+        default=True,
+        help="Whether to clean the output dir if existed. (default: %(default)s)")
+    parser.add_argument(
+        "--child_fixed_arc",
+        type=str,
+        required=True,
+        help="Architecture json file. (default: %(default)s)")
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        default="data",
+        help="Directory containing the dataset and embedding file. (default: %(default)s)")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="output",
+        help="The output directory. (default: %(default)s)")
+    parser.add_argument(
+        "--child_lr_decay_scheme",
+        type=str,
+        default="cosine",
+        help="Learning rate annealing strategy, only 'cosine' supported. (default: %(default)s)")
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=128,
+        help="Number of samples each batch for training. (default: %(default)s)")
+    parser.add_argument(
+        "--eval_batch_size",
+        type=int,
+        default=128,
+        help="Number of samples each batch for evaluation. (default: %(default)s)")
+    parser.add_argument(
+        "--class_num",
+        type=int,
+        default=5,
+        help="The number of categories. (default: %(default)s)")
+    parser.add_argument(
+        "--global_seed",
+        type=int,
+        default=1234,
+        help="Seed for reproduction. (default: %(default)s)")
+    parser.add_argument(
+        "--max_input_length",
+        type=int,
+        default=64,
+        help="The maximum length of the sentence. (default: %(default)s)")
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=10,
+        help="The number of training epochs. (default: %(default)s)")
+    parser.add_argument(
+        "--child_num_layers",
+        type=int,
+        default=24,
+        help="The layer number of the architecture. (default: %(default)s)")
+    parser.add_argument(
+        "--child_out_filters",
+        type=int,
+        default=256,
+        help="The dimension of hidden states. (default: %(default)s)")
+    parser.add_argument(
+        "--child_out_filters_scale",
+        type=int,
+        default=1,
+        help="The scale of hidden state dimension. (default: %(default)s)")
+    parser.add_argument(
+        "--child_lr_T_0",
+        type=int,
+        default=10,
+        help="The length of one cycle. (default: %(default)s)")
+    parser.add_argument(
+        "--child_lr_T_mul",
+        type=int,
+        default=2,
+        help="The multiplication factor per cycle. (default: %(default)s)")
+    parser.add_argument(
+        "--min_count",
+        type=int,
+        default=1,
+        help="The threshold to cut off low frequent words. (default: %(default)s)")
+    parser.add_argument(
+        "--train_ratio",
+        type=float,
+        default=1.0,
+        help="The sample ratio for the training set. (default: %(default)s)")
+    parser.add_argument(
+        "--valid_ratio",
+        type=float,
+        default=1.0,
+        help="The sample ratio for the dev set. (default: %(default)s)")
+    parser.add_argument(
+        "--child_grad_bound",
+        type=float,
+        default=5.0,
+        help="The threshold for gradient clipping. (default: %(default)s)")
+    parser.add_argument(
+        "--child_lr",
+        type=float,
+        default=0.02,
+        help="The initial learning rate. (default: %(default)s)")
+    parser.add_argument(
+        "--cnn_keep_prob",
+        type=float,
+        default=0.8,
+        help="Keep prob for cnn layer. (default: %(default)s)")
+    parser.add_argument(
+        "--final_output_keep_prob",
+        type=float,
+        default=1.0,
+        help="Keep prob for the last output layer. (default: %(default)s)")
+    parser.add_argument(
+        "--lstm_out_keep_prob",
+        type=float,
+        default=0.8,
+        help="Keep prob for the RNN layer. (default: %(default)s)")
+    parser.add_argument(
+        "--embed_keep_prob",
+        type=float,
+        default=0.8,
+        help="Keep prob for the embedding layer. (default: %(default)s)")
+    parser.add_argument(
+        "--attention_keep_prob",
+        type=float,
+        default=0.8,
+        help="Keep prob for the self-attention layer. (default: %(default)s)")
+    parser.add_argument(
+        "--child_l2_reg",
+        type=float,
+        default=3e-6,
+        help="Weight decay factor. (default: %(default)s)")
+    parser.add_argument(
+        "--child_lr_max",
+        type=float,
+        default=0.002,
+        help="The max learning rate. (default: %(default)s)")
+    parser.add_argument(
+        "--child_lr_min",
+        type=float,
+        default=0.001,
+        help="The min learning rate. (default: %(default)s)")
+    parser.add_argument(
+        "--child_optim_algo",
+        type=str,
+        default="adam",
+        help="Optimization algorithm. (default: %(default)s)")
+    parser.add_argument(
+        "--checkpoint_dir",
+        type=str,
+        default="best_checkpoint",
+        help="Path for saved checkpoints. (default: %(default)s)")
+    parser.add_argument(
+        "--output_type",
+        type=str,
+        default="avg",
+        help="Opertor type for the time steps reduction. (default: %(default)s)")
+    parser.add_argument(
+        "--multi_path",
+        type=distutils.util.strtobool,
+        default=False,
+        help="Search for multiple path in the architecture. (default: %(default)s)")
+    parser.add_argument(
+        "--is_binary",
+        type=distutils.util.strtobool,
+        default=False,
+        help="Binary label for sst dataset. (default: %(default)s)")
+    parser.add_argument(
+        "--is_cuda",
+        type=distutils.util.strtobool,
+        default=True,
+        help="Specify the device type. (default: %(default)s)")
+    parser.add_argument(
+        "--is_mask",
+        type=distutils.util.strtobool,
+        default=True,
+        help="Apply mask. (default: %(default)s)")
+    parser.add_argument(
+        "--fixed_seed",
+        type=distutils.util.strtobool,
+        default=True,
+        help="Fix the seed. (default: %(default)s)")
+    parser.add_argument(
+        "--load_checkpoint",
+        type=distutils.util.strtobool,
+        default=False,
+        help="Wether to load checkpoint. (default: %(default)s)")
+    parser.add_argument(
+        "--log_every",
+        type=int,
+        default=50,
+        help="How many steps to log. (default: %(default)s)")
+    parser.add_argument(
+        "--eval_every_epochs",
+        type=int,
+        default=1,
+        help="How many epochs to eval. (default: %(default)s)")
+
+    global FLAGS
+
+    FLAGS = parser.parse_args()
+
+
+def set_random_seed(seed):
+    logger.info("set random seed for data reading: {}".format(seed))
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if FLAGS.is_cuda:
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.deterministic = True
+
+
+def get_model(embedding, num_layers):
+    logger.info("num layers: {0}".format(num_layers))
+    assert FLAGS.child_fixed_arc is not None, "Architecture should be provided."
+
+    child_model = Model(
+        embedding=embedding,
+        hidden_units=FLAGS.child_out_filters_scale * FLAGS.child_out_filters,
+        num_layers=num_layers,
+        num_classes=FLAGS.class_num,
+        choose_from_k=5 if FLAGS.multi_path else 1,
+        lstm_keep_prob=FLAGS.lstm_out_keep_prob,
+        cnn_keep_prob=FLAGS.cnn_keep_prob,
+        att_keep_prob=FLAGS.attention_keep_prob,
+        att_mask=FLAGS.is_mask,
+        embed_keep_prob=FLAGS.embed_keep_prob,
+        final_output_keep_prob=FLAGS.final_output_keep_prob,
+        global_pool=FLAGS.output_type)
+
+    apply_fixed_architecture(child_model, FLAGS.child_fixed_arc)
+    return child_model
+
+
+def eval_once(child_model, device, eval_set, criterion, valid_dataloader=None, test_dataloader=None):
+    if eval_set == "test":
+        assert test_dataloader is not None
+        dataloader = test_dataloader
+    elif eval_set == "valid":
+        assert valid_dataloader is not None
+        dataloader = valid_dataloader
+    else:
+        raise NotImplementedError("Unknown eval_set '{}'".format(eval_set))
+
+    tot_acc = 0
+    tot = 0
+    losses = []
+
+    with torch.no_grad():  # save memory
+        for batch in dataloader:
+            (sent_ids, mask), labels = batch
+
+            sent_ids = sent_ids.to(device, non_blocking=True)
+            mask = mask.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
+
+            logits = child_model((sent_ids, mask))  # run
+
+            loss = criterion(logits, labels.long())
+            loss = loss.mean()
+            preds = logits.argmax(dim=1).long()
+            acc = torch.eq(preds, labels.long()).long().sum().item()
+
+            losses.append(loss)
+            tot_acc += acc
+            tot += len(labels)
+
+    losses = torch.tensor(losses)
+    loss = losses.mean()
+    if tot > 0:
+        final_acc = float(tot_acc) / tot
+    else:
+        final_acc = 0
+        logger.info("Error in calculating final_acc")
+    return final_acc, loss
+
+
+def print_user_flags(FLAGS, line_limit=80):
+    log_strings = "\n" + "-" * line_limit + "\n"
+    for flag_name in sorted(vars(FLAGS)):
+        value = "{}".format(getattr(FLAGS, flag_name))
+        log_string = flag_name
+        log_string += "." * (line_limit - len(flag_name) - len(value))
+        log_string += value
+        log_strings = log_strings + log_string
+        log_strings = log_strings + "\n"
+    log_strings += "-" * line_limit
+    logger.info(log_strings)
+
+
+def count_model_params(trainable_params):
+    num_vars = 0
+    for var in trainable_params:
+        num_vars += np.prod([dim for dim in var.size()])
+    return num_vars
+
+
+def update_lr(
+        optimizer,
+        epoch,
+        l2_reg=1e-4,
+        lr_warmup_val=None,
+        lr_init=0.1,
+        lr_decay_scheme="cosine",
+        lr_max=0.002,
+        lr_min=0.000000001,
+        lr_T_0=4,
+        lr_T_mul=1,
+        sync_replicas=False,
+        num_aggregate=None,
+        num_replicas=None):
+    if lr_decay_scheme == "cosine":
+        assert lr_max is not None, "Need lr_max to use lr_cosine"
+        assert lr_min is not None, "Need lr_min to use lr_cosine"
+        assert lr_T_0 is not None, "Need lr_T_0 to use lr_cosine"
+        assert lr_T_mul is not None, "Need lr_T_mul to use lr_cosine"
+
+        T_i = lr_T_0
+        t_epoch = epoch
+        last_reset = 0
+        while True:
+            t_epoch -= T_i
+            if t_epoch < 0:
+              break
+            last_reset += T_i
+            T_i *= lr_T_mul
+
+        T_curr = epoch - last_reset
+
+        def _update():
+            rate = T_curr / T_i * 3.1415926
+            lr = lr_min + 0.5 * (lr_max - lr_min) * (1.0 + math.cos(rate))
+            return lr
+
+        learning_rate = _update()
+    else:
+        raise ValueError("Unknown learning rate decay scheme {}".format(lr_decay_scheme))
+
+    #update lr in optimizer
+    for params_group in optimizer.param_groups:
+        params_group['lr'] = learning_rate
+    return learning_rate
+
+
+def train(device, data_path, output_dir, num_layers):
+    logger.info("Build dataloader")
+    train_dataset, valid_dataset, test_dataset, embedding = \
+        read_data_sst(data_path,
+                      FLAGS.max_input_length,
+                      FLAGS.min_count,
+                      train_ratio=FLAGS.train_ratio,
+                      valid_ratio=FLAGS.valid_ratio,
+                      is_binary=FLAGS.is_binary)
+    train_dataloader = DataLoader(train_dataset, batch_size=FLAGS.batch_size, shuffle=True, pin_memory=True)
+    test_dataloader = DataLoader(test_dataset, batch_size=FLAGS.eval_batch_size, pin_memory=True)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=FLAGS.eval_batch_size, pin_memory=True)
+
+    logger.info("Build model")
+    child_model = get_model(embedding, num_layers)
+    logger.info("Finish build model")
+
+    #for name, var in child_model.named_parameters():
+    #    logger.info(name, var.size(), var.requires_grad)  # output all params
+
+    num_vars = count_model_params(child_model.parameters())
+    logger.info("Model has {} params".format(num_vars))
+
+    for m in child_model.modules():  # initializer
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.xavier_uniform_(m.weight)
+
+    criterion = nn.CrossEntropyLoss()
+
+    # get optimizer
+    if FLAGS.child_optim_algo == "adam":
+        optimizer = optim.Adam(child_model.parameters(), eps=1e-3, weight_decay=FLAGS.child_l2_reg)  # with L2
+    else:
+        raise ValueError("Unknown optim_algo {}".format(FLAGS.child_optim_algo))
+
+    child_model.to(device)
+    criterion.to(device)
+
+    logger.info("Start training")
+    start_time = time.time()
+    step = 0
+
+    # save path
+    model_save_path = os.path.join(FLAGS.output_dir, "model.pth")
+    best_model_save_path = os.path.join(FLAGS.output_dir, "best_model.pth")
+    best_acc = 0
+    start_epoch = 0
+    if FLAGS.load_checkpoint:
+        if os.path.isfile(model_save_path):
+            checkpoint = torch.load(model_save_path, map_location = torch.device('cpu'))
+            step = checkpoint['step']
+            start_epoch = checkpoint['epoch']
+            child_model.load_state_dict(checkpoint['child_model_state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+    for epoch in range(start_epoch, FLAGS.num_epochs):
+        lr = update_lr(optimizer,
+                       epoch,
+                       l2_reg=FLAGS.child_l2_reg,
+                       lr_warmup_val=None,
+                       lr_init=FLAGS.child_lr,
+                       lr_decay_scheme=FLAGS.child_lr_decay_scheme,
+                       lr_max=FLAGS.child_lr_max,
+                       lr_min=FLAGS.child_lr_min,
+                       lr_T_0=FLAGS.child_lr_T_0,
+                       lr_T_mul=FLAGS.child_lr_T_mul)
+        child_model.train()
+        for batch in train_dataloader:
+            (sent_ids, mask), labels = batch
+
+            sent_ids = sent_ids.to(device, non_blocking=True)
+            mask = mask.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
+
+            step += 1
+
+            logits = child_model((sent_ids, mask))  # run
+
+            loss = criterion(logits, labels.long())
+            loss = loss.mean()
+            preds = logits.argmax(dim=1).long()
+            acc = torch.eq(preds, labels.long()).long().sum().item()
+
+            optimizer.zero_grad()
+            loss.backward()
+            grad_norm = 0
+            trainable_params = child_model.parameters()
+
+            assert FLAGS.child_grad_bound is not None, "Need grad_bound to clip gradients."
+            # compute the gradient norm value
+            grad_norm = nn.utils.clip_grad_norm_(trainable_params, 99999999)
+            for param in trainable_params:
+                nn.utils.clip_grad_norm_(param, FLAGS.child_grad_bound)  # clip grad
+
+            optimizer.step()
+
+            if step % FLAGS.log_every == 0:
+                curr_time = time.time()
+                log_string = ""
+                log_string += "epoch={:<6d}".format(epoch)
+                log_string += "ch_step={:<6d}".format(step)
+                log_string += " loss={:<8.6f}".format(loss)
+                log_string += " lr={:<8.4f}".format(lr)
+                log_string += " |g|={:<8.4f}".format(grad_norm)
+                log_string += " tr_acc={:<3d}/{:>3d}".format(acc, logits.size()[0])
+                log_string += " mins={:<10.2f}".format(float(curr_time - start_time) / 60)
+                logger.info(log_string)
+
+        epoch += 1
+        save_state = {
+            'step' : step,
+            'epoch' : epoch,
+            'child_model_state_dict' : child_model.state_dict(),
+            'optimizer_state_dict' : optimizer.state_dict()}
+        torch.save(save_state, model_save_path)
+        child_model.eval()
+        logger.info("Epoch {}: Eval".format(epoch))
+        eval_acc, eval_loss = eval_once(child_model, device, "test", criterion, test_dataloader=test_dataloader)
+        logger.info("ch_step={} {}_accuracy={:<6.4f} {}_loss={:<6.4f}".format(step, "test", eval_acc, "test", eval_loss))
+        if eval_acc > best_acc:
+            best_acc = eval_acc
+            logger.info("Save best model")
+            save_state = {
+                'step' : step,
+                'epoch' : epoch,
+                'child_model_state_dict' : child_model.state_dict(),
+                'optimizer_state_dict' : optimizer.state_dict()}
+            torch.save(save_state, best_model_save_path)
+
+    return eval_acc
+
+
+def main():
+    parse_args()
+    if not os.path.isdir(FLAGS.output_dir):
+        logger.info("Path {} does not exist. Creating.".format(FLAGS.output_dir))
+        os.makedirs(FLAGS.output_dir)
+    elif FLAGS.reset_output_dir:
+        logger.info("Path {} exists. Remove and remake.".format(FLAGS.output_dir))
+        shutil.rmtree(FLAGS.output_dir, ignore_errors=True)
+        os.makedirs(FLAGS.output_dir)
+
+    print_user_flags(FLAGS)
+
+    if FLAGS.fixed_seed:
+        set_random_seed(FLAGS.global_seed)
+
+    device = torch.device("cuda" if FLAGS.is_cuda else "cpu")
+    train(device, FLAGS.data_path, FLAGS.output_dir, FLAGS.child_num_layers)
+
+
+if __name__ == "__main__":
+  main()

examples/nas/textnas/run_retrain.sh

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+export PYTHONPATH="$(pwd)"
+export CUDA_VISIBLE_DEVICES=0
+
+python -u retrain.py \
+  --train_ratio=1.0 \
+  --valid_ratio=1.0 \
+  --min_count=1 \
+  --is_mask=True \
+  --is_binary=True \
+  --child_lr_decay_scheme="cosine" \
+  --data_path="data" \
+  --class_num=2 \
+  --child_optim_algo="adam" \
+  --output_dir="output_sst2" \
+  --global_seed=1234 \
+  --max_input_length=64 \
+  --batch_size=128 \
+  --eval_batch_size=128 \
+  --num_epochs=10 \
+  --log_every=50 \
+  --eval_every_epochs=1 \
+  --child_num_layers=24 \
+  --child_out_filters=256 \
+  --child_l2_reg=1e-6 \
+  --cnn_keep_prob=0.8 \
+  --final_output_keep_prob=1.0 \
+  --embed_keep_prob=0.8 \
+  --lstm_out_keep_prob=0.8 \
+  --attention_keep_prob=0.8 \
+  --child_lr=0.02 \
+  --child_lr_max=0.002 \
+  --child_lr_min=5e-6 \
+  --child_lr_T_0=10 \
+  --child_lr_T_mul=2 \
+  --multi_path=True \
+  --child_fixed_arc="./checkpoints/architecture_00.json" \
+  --fixed_seed=True \
+  "$@"

examples/trials/mnist-pbt-tuner-pytorch/.gitignore

+tmp

examples/trials/mnist-pbt-tuner-pytorch/config.yml

+authorName: default
+experimentName: example_mnist_pbt_tuner_pytorch
+trialConcurrency: 3
+maxExecDuration: 2h
+maxTrialNum: 100
+#choice: local, remote, pai
+trainingServicePlatform: local
+searchSpacePath: search_space.json
+#choice: true, false
+useAnnotation: false
+tuner:
+#  codeDir: ~/nni/src/sdk/pynni/nni/pbt_tuner
+#  classFileName: pbt_tuner.py
+#  className: PBTTuner
+  builtinTunerName: PBTTuner
+  classArgs:
+    #choice: maximize, minimize
+    optimize_mode: maximize
+trial:
+  command: python3 mnist.py
+  codeDir: .
+  gpuNum: 1

examples/trials/mnist-pbt-tuner-pytorch/mnist.py

+import argparse
+import logging
+
+import os
+import nni
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+
+
+logger = logging.getLogger('mnist_pbt_tuner_pytorch_AutoML')
+
+class Net(nn.Module):
+    def __init__(self, hidden_size):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5, 1)
+        self.conv2 = nn.Conv2d(20, 50, 5, 1)
+        self.fc1 = nn.Linear(4*4*50, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, 10)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.max_pool2d(x, 2, 2)
+        x = F.relu(self.conv2(x))
+        x = F.max_pool2d(x, 2, 2)
+        x = x.view(-1, 4*4*50)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def train(args, model, device, train_loader, optimizer, epoch):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args['log_interval'] == 0:
+            logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss.item()))
+
+
+def test(args, model, device, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            # sum up batch loss
+            test_loss += F.nll_loss(output, target, reduction='sum').item()
+            # get the index of the max log-probability
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+
+    accuracy = 100. * correct / len(test_loader.dataset)
+
+    logger.info('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset), accuracy))
+
+    return accuracy
+
+
+def save_checkpoint(model, checkpoint_path):
+    torch.save(model.state_dict(), checkpoint_path)
+
+
+def load_checkpoint(checkpoint_path):
+    model_state_dict = torch.load(checkpoint_path)
+    return model_state_dict
+
+
+def main(args):
+    use_cuda = not args['no_cuda'] and torch.cuda.is_available()
+
+    torch.manual_seed(args['seed'])
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
+
+    data_dir = os.path.join(args['data_dir'], nni.get_trial_id())
+
+    train_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(data_dir, train=True, download=True,
+                       transform=transforms.Compose([
+                           transforms.ToTensor(),
+                           transforms.Normalize((0.1307,), (0.3081,))
+                       ])),
+        batch_size=args['batch_size'], shuffle=True, **kwargs)
+    test_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(data_dir, train=False, transform=transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307,), (0.3081,))
+        ])),
+        batch_size=1000, shuffle=True, **kwargs)
+
+    hidden_size = args['hidden_size']
+
+    model = Net(hidden_size=hidden_size).to(device)
+
+    save_checkpoint_dir = args['save_checkpoint_dir']
+    save_checkpoint_path = os.path.join(save_checkpoint_dir, 'model.pth')
+    load_checkpoint_path = os.path.join(args['load_checkpoint_dir'], 'model.pth')
+
+    if os.path.isfile(load_checkpoint_path):
+        model_state_dict = load_checkpoint(load_checkpoint_path)
+        logger.info("test : " + load_checkpoint_path)
+        logger.info(type(model_state_dict))
+        model.load_state_dict(model_state_dict)
+
+    optimizer = optim.SGD(model.parameters(), lr=args['lr'],
+                          momentum=args['momentum'])
+
+    #epoch is perturbation interval
+    for epoch in range(1, args['epochs'] + 1):
+        train(args, model, device, train_loader, optimizer, epoch)
+        test_acc = test(args, model, device, test_loader)
+
+        if epoch < args['epochs']:
+            # report intermediate result
+            nni.report_intermediate_result(test_acc)
+            logger.debug('test accuracy %g', test_acc)
+            logger.debug('Pipe send intermediate result done.')
+        else:
+            # report final result
+            nni.report_final_result(test_acc)
+            logger.debug('Final result is %g', test_acc)
+            logger.debug('Send final result done.')
+
+    if not os.path.exists(save_checkpoint_dir):
+        os.makedirs(save_checkpoint_dir)
+    save_checkpoint(model, save_checkpoint_path)
+
+
+def get_params():
+    # Training settings
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+    parser.add_argument("--data_dir", type=str,
+                        default='./tmp/pytorch/mnist/input_data', help="data directory")
+    parser.add_argument('--batch_size', type=int, default=64, metavar='N',
+                        help='input batch size for training (default: 64)')
+    parser.add_argument("--hidden_size", type=int, default=512, metavar='N',
+                        help='hidden layer size (default: 512)')
+    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
+                        help='learning rate (default: 0.01)')
+    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
+                        help='SGD momentum (default: 0.5)')
+    parser.add_argument('--epochs', type=int, default=10, metavar='N',
+                        help='number of epochs to train (default: 10)')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--no_cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--log_interval', type=int, default=1000, metavar='N',
+                        help='how many batches to wait before logging training status')
+
+    parser.add_argument('--save_checkpoint_dir', type=str,
+                        help='where to save checkpoint of this trial')
+    parser.add_argument('--load_checkpoint_dir', type=str,
+                        help='where to load the model')
+
+
+    args, _ = parser.parse_known_args()
+    return args
+
+
+if __name__ == '__main__':
+    try:
+        # get parameters form tuner
+        tuner_params = nni.get_next_parameter()
+        logger.debug(tuner_params)
+        params = vars(get_params())
+        params.update(tuner_params)
+        main(params)
+    except Exception as exception:
+        logger.exception(exception)
+        raise
\ No newline at end of file

examples/trials/mnist-pbt-tuner-pytorch/search_space.json

+{
+    "batch_size": {"_type":"choice", "_value": [16, 32, 64, 128]},
+    "hidden_size":{"_type":"choice","_value":[128, 256, 512, 1024]},
+    "lr":{"_type":"choice","_value":[0.0001, 0.001, 0.01, 0.1]},
+    "momentum":{"_type":"uniform","_value":[0, 1]}
+}

examples/trials/mnist-pytorch/mnist.py

@@ -15,14 +15,6 @@ import torch.nn.functional as F
 import torch.optim as optim
 from torchvision import datasets, transforms
 
-# Temporary patch this example until the MNIST dataset download issue get resolved
-# https://github.com/pytorch/vision/issues/1938
-import urllib
-
-opener = urllib.request.build_opener()
-opener.addheaders = [('User-agent', 'Mozilla/5.0')]
-urllib.request.install_opener(opener)
-
 logger = logging.getLogger('mnist_AutoML')
 
 
@@ -48,6 +40,8 @@ class Net(nn.Module):
 def train(args, model, device, train_loader, optimizer, epoch):
     model.train()
     for batch_idx, (data, target) in enumerate(train_loader):
+        if (args['batch_num'] is not None) and batch_idx >= args['batch_num']:
+            break
         data, target = data.to(device), target.to(device)
         optimizer.zero_grad()
         output = model(data)
@@ -119,12 +113,11 @@ def main(args):
         train(args, model, device, train_loader, optimizer, epoch)
         test_acc = test(args, model, device, test_loader)
 
-        if epoch < args['epochs']:
         # report intermediate result
         nni.report_intermediate_result(test_acc)
         logger.debug('test accuracy %g', test_acc)
         logger.debug('Pipe send intermediate result done.')
-        else:
+
     # report final result
     nni.report_final_result(test_acc)
     logger.debug('Final result is %g', test_acc)
@@ -138,6 +131,7 @@ def get_params():
                         default='/tmp/pytorch/mnist/input_data', help="data directory")
     parser.add_argument('--batch_size', type=int, default=64, metavar='N',
                         help='input batch size for training (default: 64)')
+    parser.add_argument("--batch_num", type=int, default=None)
     parser.add_argument("--hidden_size", type=int, default=512, metavar='N',
                         help='hidden layer size (default: 512)')
     parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
@@ -165,6 +159,7 @@ if __name__ == '__main__':
         logger.debug(tuner_params)
         params = vars(get_params())
         params.update(tuner_params)
+        print(params)
         main(params)
     except Exception as exception:
         logger.exception(exception)

src/nni_manager/rest_server/restValidationSchemas.ts

@@ -178,7 +178,7 @@ export namespace ValidationSchemas {
                 gpuIndices: joi.string()
             }),
             tuner: joi.object({
-                builtinTunerName: joi.string().valid('TPE', 'Random', 'Anneal', 'Evolution', 'SMAC', 'BatchTuner', 'GridSearch', 'NetworkMorphism', 'MetisTuner', 'GPTuner', 'PPOTuner'),
+                builtinTunerName: joi.string().valid('TPE', 'Random', 'Anneal', 'Evolution', 'SMAC', 'BatchTuner', 'GridSearch', 'NetworkMorphism', 'MetisTuner', 'GPTuner', 'PPOTuner', 'PBTTuner'),
                 codeDir: joi.string(),
                 classFileName: joi.string(),
                 className: joi.string(),

src/sdk/pynni/nni/compression/speedup/torch/compressor.py

@@ -229,42 +229,39 @@ class ModelSpeedup:
         list
             a list of scope name of all the leaf modules
         """
-        pieces = [] # each element is a dict
+        class SNode:
+            def __init__(self, name):
+                self.sname = name
+                self.childs = {}
+
+        root = None
         for node in graph.nodes():
             scope_name = node.scopeName()
             if scope_name == '':
                 continue
             segs = scope_name.split('/')
-            segs_len = len(segs)
-            # increase the length of `pieces` if not enough
-            for _ in range(segs_len - len(pieces)):
-                pieces.append({})
-            # process internal segments of the scope name
-            # 'L' means leaf segment
-            # 'I' means internal segment
-            # internal segment can replace leaf segment at the same position of `pieces`
-            for i, seg in enumerate(segs[:-1]):
-                seg_name_dict = pieces[i]
-                if seg in seg_name_dict:
-                    if seg_name_dict[seg][0] == 'L':
-                        seg_name_dict[seg] = ('I', node)
+            if root is None:
+                root = SNode(segs[0])
+            curr = root
+            for seg in segs[1:]:
+                if not seg in curr.childs:
+                    curr.childs[seg] = SNode(seg)
+                curr = curr.childs[seg]
+
+        leaf_nodes = []
+        def traverse_tree(node, scope_name):
+            if scope_name == '':
+                sn = node.sname
+            else:
+                sn = scope_name + '/' + node.sname
+            if not node.childs:
+                if node.sname[-1] == ']':
+                    leaf_nodes.append(sn)
             else:
-                    seg_name_dict[seg] = ('I', node)
-            # process the leaf segment of the scope name
-            last_segs_dict = pieces[len(segs) - 1]
-            if not segs[-1] in last_segs_dict:
-                last_segs_dict[segs[-1]] = ('L', node)
-        # traverse `pieces` to obtain all the leaf modules which are labeled with 'L'
-        leaf_modules = []
-        for piece in pieces:
-            for _, value in piece.items():
-                if value[0] == 'L':
-                    assert value[1].scopeName() not in leaf_modules
-                    # if this is a leaf module, the last segment of its scope name
-                    # must be in pattern `xxx[xxx]`
-                    if value[1].scopeName()[-1] == ']':
-                        leaf_modules.append(value[1].scopeName())
-        return leaf_modules
+                for key in node.childs:
+                    traverse_tree(node.childs[key], sn)
+        traverse_tree(root, '')
+        return leaf_nodes
 
     def _build_graph(self):
         """

src/sdk/pynni/nni/constants.py

@@ -15,7 +15,8 @@ ModuleName = {
     'Curvefitting': 'nni.curvefitting_assessor.curvefitting_assessor',
     'MetisTuner': 'nni.metis_tuner.metis_tuner',
     'GPTuner': 'nni.gp_tuner.gp_tuner',
-    'PPOTuner': 'nni.ppo_tuner.ppo_tuner'
+    'PPOTuner': 'nni.ppo_tuner.ppo_tuner',
+    'PBTTuner': 'nni.pbt_tuner.pbt_tuner'
 }
 
 ClassName = {
@@ -30,6 +31,7 @@ ClassName = {
     'MetisTuner':'MetisTuner',
     'GPTuner':'GPTuner',
     'PPOTuner': 'PPOTuner',
+    'PBTTuner': 'PBTTuner',
 
     'Medianstop': 'MedianstopAssessor',
     'Curvefitting': 'CurvefittingAssessor'

src/sdk/pynni/nni/evolution_tuner/evolution_tuner.py

@@ -10,97 +10,8 @@ import random
 
 import numpy as np
 from nni.tuner import Tuner
-from nni.utils import NodeType, OptimizeMode, extract_scalar_reward, split_index
+from nni.utils import OptimizeMode, extract_scalar_reward, split_index, json2parameter, json2space
 
-import nni.parameter_expressions as parameter_expressions
-
-
-def json2space(x, oldy=None, name=NodeType.ROOT):
-    """
-    Change search space from json format to hyperopt format
-
-    """
-    y = list()
-    if isinstance(x, dict):
-        if NodeType.TYPE in x.keys():
-            _type = x[NodeType.TYPE]
-            name = name + '-' + _type
-            if _type == 'choice':
-                if oldy is not None:
-                    _index = oldy[NodeType.INDEX]
-                    y += json2space(x[NodeType.VALUE][_index],
-                                    oldy[NodeType.VALUE], name=name+'[%d]' % _index)
-                else:
-                    y += json2space(x[NodeType.VALUE], None, name=name)
-            y.append(name)
-        else:
-            for key in x.keys():
-                y += json2space(x[key], oldy[key] if oldy else None, name+"[%s]" % str(key))
-    elif isinstance(x, list):
-        for i, x_i in enumerate(x):
-            if isinstance(x_i, dict):
-                if NodeType.NAME not in x_i.keys():
-                    raise RuntimeError('\'_name\' key is not found in this nested search space.')
-            y += json2space(x_i, oldy[i] if oldy else None, name + "[%d]" % i)
-    return y
-
-def json2parameter(x, is_rand, random_state, oldy=None, Rand=False, name=NodeType.ROOT):
-    """
-    Json to pramaters.
-
-    """
-    if isinstance(x, dict):
-        if NodeType.TYPE in x.keys():
-            _type = x[NodeType.TYPE]
-            _value = x[NodeType.VALUE]
-            name = name + '-' + _type
-            Rand |= is_rand[name]
-            if Rand is True:
-                if _type == 'choice':
-                    _index = random_state.randint(len(_value))
-                    y = {
-                        NodeType.INDEX: _index,
-                        NodeType.VALUE: json2parameter(
-                            x[NodeType.VALUE][_index],
-                            is_rand,
-                            random_state,
-                            None,
-                            Rand,
-                            name=name+"[%d]" % _index
-                        )
-                    }
-                else:
-                    y = getattr(parameter_expressions, _type)(*(_value + [random_state]))
-            else:
-                y = copy.deepcopy(oldy)
-        else:
-            y = dict()
-            for key in x.keys():
-                y[key] = json2parameter(
-                    x[key],
-                    is_rand,
-                    random_state,
-                    oldy[key] if oldy else None,
-                    Rand,
-                    name + "[%s]" % str(key)
-                )
-    elif isinstance(x, list):
-        y = list()
-        for i, x_i in enumerate(x):
-            if isinstance(x_i, dict):
-                if NodeType.NAME not in x_i.keys():
-                    raise RuntimeError('\'_name\' key is not found in this nested search space.')
-            y.append(json2parameter(
-                x_i,
-                is_rand,
-                random_state,
-                oldy[i] if oldy else None,
-                Rand,
-                name + "[%d]" % i
-            ))
-    else:
-        y = copy.deepcopy(x)
-    return y
 
 class Individual:
     """

src/sdk/pynni/nni/evolution_tuner/test_evolution_tuner.py

@@ -9,7 +9,7 @@ import numpy as np
 
 from unittest import TestCase, main
 
-from nni.evolution_tuner.evolution_tuner import json2space, json2parameter
+from nni.utils import json2space, json2parameter
 
 
 class EvolutionTunerTestCase(TestCase):

src/sdk/pynni/nni/pbt_tuner/pbt_tuner.py

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import copy
+import logging
+import os
+import numpy as np
+
+import nni
+from nni.tuner import Tuner
+from nni.utils import OptimizeMode, extract_scalar_reward, split_index, json2parameter, json2space
+
+
+logger = logging.getLogger('pbt_tuner_AutoML')
+
+
+def exploit_and_explore(bot_trial_info, top_trial_info, factors, epoch, search_space):
+    """
+    Replace checkpoint of bot_trial with top, and perturb hyperparameters
+
+    Parameters
+    ----------
+    bot_trial_info : TrialInfo
+        bottom model whose parameters should be replaced
+    top_trial_info : TrialInfo
+        better model
+    factors : float
+        factors for perturbation
+    epoch : int
+        step of PBTTuner
+    search_space : dict
+        search_space to keep perturbed hyperparameters in range
+    """
+    bot_checkpoint_dir = bot_trial_info.checkpoint_dir
+    top_hyper_parameters = top_trial_info.hyper_parameters
+    hyper_parameters = copy.deepcopy(top_hyper_parameters)
+    # TODO think about different type of hyperparameters for 1.perturbation 2.within search space
+    for key in hyper_parameters.keys():
+        if key == 'load_checkpoint_dir':
+            hyper_parameters[key] = hyper_parameters['save_checkpoint_dir']
+        elif key == 'save_checkpoint_dir':
+            hyper_parameters[key] = os.path.join(bot_checkpoint_dir, str(epoch))
+        elif isinstance(hyper_parameters[key], float):
+            perturb = np.random.choice(factors)
+            val = hyper_parameters[key] * perturb
+            lb, ub = search_space[key]["_value"][:2]
+            if search_space[key]["_type"] in ("uniform", "normal"):
+                val = np.clip(val, lb, ub).item()
+                hyper_parameters[key] = val
+        else:
+            continue
+    bot_trial_info.hyper_parameters = hyper_parameters
+    bot_trial_info.clean_id()
+
+
+class TrialInfo:
+    """
+    Information of each trial, refresh for each epoch
+
+    """
+
+    def __init__(self, checkpoint_dir=None, hyper_parameters=None, parameter_id=None, score=None):
+        self.checkpoint_dir = checkpoint_dir
+        self.hyper_parameters = hyper_parameters
+        self.parameter_id = parameter_id
+        self.score = score
+
+    def clean_id(self):
+        self.parameter_id = None
+
+
+class PBTTuner(Tuner):
+    def __init__(self, optimize_mode="maximize", all_checkpoint_dir=None, population_size=10, factors=(1.2, 0.8), fraction=0.2):
+        """
+        Initialization
+
+        Parameters
+        ----------
+        optimize_mode : str
+            maximize or minimize
+        all_checkpoint_dir : str
+            directory to store training model checkpoint
+        population_size : int
+            number of trials for each epoch
+        factors : tuple
+            factors for perturbation
+        fraction : float
+            fraction for selecting bottom and top trials
+        """
+        self.optimize_mode = OptimizeMode(optimize_mode)
+        if all_checkpoint_dir is None:
+            all_checkpoint_dir = os.getenv('NNI_CHECKPOINT_DIRECTORY')
+            logger.info("Checkpoint dir is set to %s by default.", all_checkpoint_dir)
+        self.all_checkpoint_dir = all_checkpoint_dir
+        self.population_size = population_size
+        self.factors = factors
+        self.fraction = fraction
+        # defined in trial code
+        #self.perturbation_interval = perturbation_interval
+
+        self.population = None
+        self.pos = -1
+        self.param_ids = []
+        self.running = {}
+        self.finished = []
+        self.credit = 0
+        self.finished_trials = 0
+        self.epoch = 0
+
+        self.searchspace_json = None
+        self.space = None
+
+        self.send_trial_callback = None
+
+        logger.info('PBT tuner initialization')
+
+    def update_search_space(self, search_space):
+        """
+        Get search space
+
+        Parameters
+        ----------
+        search_space : dict
+            Search space
+        """
+        logger.info('Update search space %s', search_space)
+        self.searchspace_json = search_space
+        self.space = json2space(self.searchspace_json)
+
+        self.random_state = np.random.RandomState()
+        self.population = []
+        is_rand = dict()
+
+        for item in self.space:
+            is_rand[item] = True
+
+        for i in range(self.population_size):
+            hyper_parameters = json2parameter(
+                self.searchspace_json, is_rand, self.random_state)
+            checkpoint_dir = os.path.join(self.all_checkpoint_dir, str(i))
+            hyper_parameters['load_checkpoint_dir'] = os.path.join(checkpoint_dir, str(self.epoch))
+            hyper_parameters['save_checkpoint_dir'] = os.path.join(checkpoint_dir, str(self.epoch))
+            self.population.append(TrialInfo(checkpoint_dir=checkpoint_dir, hyper_parameters=hyper_parameters))
+
+    def generate_multiple_parameters(self, parameter_id_list, **kwargs):
+        """
+        Returns multiple sets of trial (hyper-)parameters, as iterable of serializable objects.
+
+        Parameters
+        ----------
+        parameter_id_list : list of int
+            Unique identifiers for each set of requested hyper-parameters.
+            These will later be used in :meth:`receive_trial_result`.
+        **kwargs
+            Used for send_trial_callback.
+
+        Returns
+        -------
+        list
+            A list of newly generated configurations
+        """
+        result = []
+        self.send_trial_callback = kwargs['st_callback']
+        for parameter_id in parameter_id_list:
+            had_exception = False
+            try:
+                logger.debug("generating param for %s", parameter_id)
+                res = self.generate_parameters(parameter_id, **kwargs)
+            except nni.NoMoreTrialError:
+                had_exception = True
+            if not had_exception:
+                result.append(res)
+        return result
+
+    def generate_parameters(self, parameter_id, **kwargs):
+        """
+        Generate parameters, if no trial configration for now, self.credit plus 1 to send the config later
+
+        Parameters
+        ----------
+        parameter_id : int
+            Unique identifier for requested hyper-parameters.
+            This will later be used in :meth:`receive_trial_result`.
+        **kwargs
+            Not used
+
+        Returns
+        -------
+        dict
+            One newly generated configuration
+
+        """
+        if self.pos == self.population_size - 1:
+            logger.debug('Credit added by one in parameters request')
+            self.credit += 1
+            self.param_ids.append(parameter_id)
+            raise nni.NoMoreTrialError('No more parameters now.')
+        self.pos += 1
+        trial_info = self.population[self.pos]
+        trial_info.parameter_id = parameter_id
+        self.running[parameter_id] = trial_info
+        logger.info('Generate parameter : %s', trial_info.hyper_parameters)
+        return split_index(trial_info.hyper_parameters)
+
+    def receive_trial_result(self, parameter_id, parameters, value, **kwargs):
+        """
+        Receive trial's result. if the number of finished trials equals ``self.population_size``, start the next epoch to
+        train the model.
+
+        Parameters
+        ----------
+        parameter_id : int
+            Unique identifier of used hyper-parameters, same with :meth:`generate_parameters`.
+        parameters : dict
+            Hyper-parameters generated by :meth:`generate_parameters`.
+        value : dict
+            Result from trial (the return value of :func:`nni.report_final_result`).
+        """
+        logger.info('Get one trial result, id = %d, value = %s', parameter_id, value)
+        value = extract_scalar_reward(value)
+        if self.optimize_mode == OptimizeMode.Minimize:
+            value = -value
+        trial_info = self.running.pop(parameter_id, None)
+        trial_info.score = value
+        self.finished.append(trial_info)
+        self.finished_trials += 1
+        if self.finished_trials == self.population_size:
+            logger.info('Proceeding to next epoch')
+            self.epoch += 1
+            self.population = []
+            self.pos = -1
+            self.running = {}
+            #exploit and explore
+            self.finished = sorted(self.finished, key=lambda x: x.score, reverse=True)
+            cutoff = int(np.ceil(self.fraction * len(self.finished)))
+            tops = self.finished[:cutoff]
+            bottoms = self.finished[self.finished_trials - cutoff:]
+            for bottom in bottoms:
+                top = np.random.choice(tops)
+                exploit_and_explore(bottom, top, self.factors, self.epoch, self.searchspace_json)
+            for trial in self.finished:
+                if trial not in bottoms:
+                    trial.clean_id()
+                    trial.hyper_parameters['load_checkpoint_dir'] = trial.hyper_parameters['save_checkpoint_dir']
+                    trial.hyper_parameters['save_checkpoint_dir'] = os.path.join(trial.checkpoint_dir, str(self.epoch))
+            self.finished_trials = 0
+            for _ in range(self.population_size):
+                trial_info = self.finished.pop()
+                self.population.append(trial_info)
+            while self.credit > 0 and self.pos + 1 < len(self.population):
+                self.credit -= 1
+                self.pos += 1
+                parameter_id = self.param_ids.pop()
+                trial_info = self.population[self.pos]
+                trial_info.parameter_id = parameter_id
+                self.running[parameter_id] = trial_info
+                self.send_trial_callback(parameter_id, split_index(trial_info.hyper_parameters))
+
+    def import_data(self, data):
+        pass

src/sdk/pynni/nni/utils.py

@@ -2,13 +2,16 @@
 # Licensed under the MIT license.
 
 import os
+import copy
 import functools
 from enum import Enum, unique
 import json_tricks
 
+from . import parameter_expressions
 from .common import init_logger
 from .env_vars import dispatcher_env_vars
 
+
 to_json = functools.partial(json_tricks.dumps, allow_nan=True)
 
 @unique
@@ -124,3 +127,92 @@ def init_dispatcher_logger():
     if dispatcher_env_vars.NNI_LOG_DIRECTORY is not None:
         logger_file_path = os.path.join(dispatcher_env_vars.NNI_LOG_DIRECTORY, logger_file_path)
     init_logger(logger_file_path, dispatcher_env_vars.NNI_LOG_LEVEL)
+
+
+def json2space(x, oldy=None, name=NodeType.ROOT):
+    """
+    Change search space from json format to hyperopt format
+
+    """
+    y = list()
+    if isinstance(x, dict):
+        if NodeType.TYPE in x.keys():
+            _type = x[NodeType.TYPE]
+            name = name + '-' + _type
+            if _type == 'choice':
+                if oldy is not None:
+                    _index = oldy[NodeType.INDEX]
+                    y += json2space(x[NodeType.VALUE][_index],
+                                    oldy[NodeType.VALUE], name=name+'[%d]' % _index)
+                else:
+                    y += json2space(x[NodeType.VALUE], None, name=name)
+            y.append(name)
+        else:
+            for key in x.keys():
+                y += json2space(x[key], oldy[key] if oldy else None, name+"[%s]" % str(key))
+    elif isinstance(x, list):
+        for i, x_i in enumerate(x):
+            if isinstance(x_i, dict):
+                if NodeType.NAME not in x_i.keys():
+                    raise RuntimeError('\'_name\' key is not found in this nested search space.')
+            y += json2space(x_i, oldy[i] if oldy else None, name + "[%d]" % i)
+    return y
+
+
+def json2parameter(x, is_rand, random_state, oldy=None, Rand=False, name=NodeType.ROOT):
+    """
+    Json to pramaters.
+
+    """
+    if isinstance(x, dict):
+        if NodeType.TYPE in x.keys():
+            _type = x[NodeType.TYPE]
+            _value = x[NodeType.VALUE]
+            name = name + '-' + _type
+            Rand |= is_rand[name]
+            if Rand is True:
+                if _type == 'choice':
+                    _index = random_state.randint(len(_value))
+                    y = {
+                        NodeType.INDEX: _index,
+                        NodeType.VALUE: json2parameter(
+                            x[NodeType.VALUE][_index],
+                            is_rand,
+                            random_state,
+                            None,
+                            Rand,
+                            name=name+"[%d]" % _index
+                        )
+                    }
+                else:
+                    y = getattr(parameter_expressions, _type)(*(_value + [random_state]))
+            else:
+                y = copy.deepcopy(oldy)
+        else:
+            y = dict()
+            for key in x.keys():
+                y[key] = json2parameter(
+                    x[key],
+                    is_rand,
+                    random_state,
+                    oldy[key] if oldy else None,
+                    Rand,
+                    name + "[%s]" % str(key)
+                )
+    elif isinstance(x, list):
+        y = list()
+        for i, x_i in enumerate(x):
+            if isinstance(x_i, dict):
+                if NodeType.NAME not in x_i.keys():
+                    raise RuntimeError('\'_name\' key is not found in this nested search space.')
+            y.append(json2parameter(
+                x_i,
+                is_rand,
+                random_state,
+                oldy[i] if oldy else None,
+                Rand,
+                name + "[%d]" % i
+            ))
+    else:
+        y = copy.deepcopy(x)
+    return y

src/sdk/pynni/tests/test_builtin_tuners.py

@@ -8,6 +8,7 @@ import os
 import random
 import shutil
 import sys
+from collections import deque
 from unittest import TestCase, main
 
 from nni.batch_tuner.batch_tuner import BatchTuner
@@ -16,6 +17,8 @@ from nni.gp_tuner.gp_tuner import GPTuner
 from nni.gridsearch_tuner.gridsearch_tuner import GridSearchTuner
 from nni.hyperopt_tuner.hyperopt_tuner import HyperoptTuner
 from nni.metis_tuner.metis_tuner import MetisTuner
+from nni.msg_dispatcher import _pack_parameter, MsgDispatcher
+from nni.pbt_tuner.pbt_tuner import PBTTuner
 
 try:
     from nni.smac_tuner.smac_tuner import SMACTuner
@@ -23,6 +26,7 @@ except ImportError:
     assert sys.platform == "win32"
 from nni.tuner import Tuner
 
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger('test_tuner')
 
@@ -44,18 +48,29 @@ class BuiltinTunersTestCase(TestCase):
         self.params_each_round = 50
         self.exhaustive = False
 
+    def send_trial_callback(self, param_queue):
+        def receive(*args):
+            param_queue.append(tuple(args))
+        return receive
+
     def search_space_test_one(self, tuner_factory, search_space):
         tuner = tuner_factory()
         self.assertIsInstance(tuner, Tuner)
         tuner.update_search_space(search_space)
 
         for i in range(self.test_round):
+            queue = deque()
             parameters = tuner.generate_multiple_parameters(list(range(i * self.params_each_round,
-                                                                       (i + 1) * self.params_each_round)))
+                                                                       (i + 1) * self.params_each_round)),
+                                                            st_callback=self.send_trial_callback(queue))
             logger.debug(parameters)
             self.check_range(parameters, search_space)
             for k in range(min(len(parameters), self.params_each_round)):
                 tuner.receive_trial_result(self.params_each_round * i + k, parameters[k], random.uniform(-100, 100))
+            while queue:
+                id_, params = queue.popleft()
+                self.check_range([params], search_space)
+                tuner.receive_trial_result(id_, params, random.uniform(-100, 100))
             if not parameters and not self.exhaustive:
                 raise ValueError("No parameters generated")
 
@@ -65,6 +80,9 @@ class BuiltinTunersTestCase(TestCase):
             if self._testMethodName == "test_batch":
                 param = {list(search_space.keys())[0]: param}
             for k, v in param.items():
+                if k == "load_checkpoint_dir" or k == "save_checkpoint_dir":
+                    self.assertIsInstance(v, str)
+                    continue
                 if k.startswith("_mutable_layer"):
                     _, block, layer, choice = k.split("/")
                     cand = search_space[block]["_value"][layer].get(choice)
@@ -270,6 +288,16 @@ class BuiltinTunersTestCase(TestCase):
     def test_ppo(self):
         pass
 
+    def test_pbt(self):
+        self.search_space_test_all(lambda: PBTTuner(
+            all_checkpoint_dir=os.path.expanduser("~/nni/checkpoint/test/"),
+            population_size=12
+        ))
+        self.search_space_test_all(lambda: PBTTuner(
+            all_checkpoint_dir=os.path.expanduser("~/nni/checkpoint/test/"),
+            population_size=100
+        ))
+
     def tearDown(self):
         file_list = glob.glob("smac3*") + ["param_config_space.pcs", "scenario.txt", "model_path"]
         for file in file_list:

src/webui/yarn.lock

@@ -1518,8 +1518,8 @@ acorn-jsx@^5.0.0:
   resolved "https://registry.yarnpkg.com/acorn-jsx/-/acorn-jsx-5.1.0.tgz#294adb71b57398b0680015f0a38c563ee1db5384"
 
 acorn@^6.0.5, acorn@^6.0.7:
-  version "6.4.0"
-  resolved "https://registry.yarnpkg.com/acorn/-/acorn-6.4.0.tgz#b659d2ffbafa24baf5db1cdbb2c94a983ecd2784"
+  version "6.4.1"
+  resolved "https://registry.yarnpkg.com/acorn/-/acorn-6.4.1.tgz#531e58ba3f51b9dacb9a6646ca4debf5b14ca474"
 
 address@1.1.2, address@^1.0.1:
   version "1.1.2"

test/config/examples/mnist-pytorch.yml

 authorName: nni
 experimentName: default_test
 maxExecDuration: 15m
-maxTrialNum: 4
+maxTrialNum: 2
 trialConcurrency: 2
 searchSpacePath: ./mnist_pytorch_search_space.json
 
@@ -13,7 +13,7 @@ assessor:
     optimize_mode: maximize
 trial:
   codeDir: ../../../examples/trials/mnist-pytorch
-  command: python3 mnist.py  --epochs 1
+  command: python3 mnist.py  --epochs 1 --batch_num 10
   gpuNum: 0
 
 useAnnotation: false