search.py

import time
import json
import random
import logging
import argparse
import numpy as np
import torch
import torchvision.transforms as transforms
import torchvision.datasets as datasets

import nni
from nn_meter import load_latency_predictor
import nni.retiarii.nn.pytorch as nn
import nni.retiarii.strategy as strategy
from nni.retiarii.evaluator.functional import FunctionalEvaluator
from nni.retiarii.utils import original_state_dict_hooks
from nni.retiarii.oneshot.pytorch.utils import AverageMeterGroup
from nni.retiarii.experiment.pytorch import RetiariiExeConfig, RetiariiExperiment

from network import ShuffleNetV2OneShot, load_and_parse_state_dict
from utils import CrossEntropyLabelSmooth, accuracy, ToBGRTensor, get_archchoice_by_model

logger = logging.getLogger("nni.spos.search")


def retrain_bn(model, criterion, max_iters, log_freq, loader):
    with torch.no_grad():
        logger.info("Clear BN statistics...")
        for m in model.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.running_mean = torch.zeros_like(m.running_mean)
                m.running_var = torch.ones_like(m.running_var)

        logger.info("Train BN with training set (BN sanitize)...")
        model.train()
        meters = AverageMeterGroup()
        for step in range(max_iters):
            inputs, targets = next(iter(loader))
            inputs, targets = inputs.to('cuda'), targets.to('cuda')
            logits = model(inputs)
            loss = criterion(logits, targets)
            metrics = accuracy(logits, targets)
            metrics["loss"] = loss.item()
            meters.update(metrics)
            if step % log_freq == 0 or step + 1 == max_iters:
                logger.info("Train Step [%d/%d] %s", step + 1, max_iters, meters)


def test_acc(model, criterion, log_freq, loader):
    logger.info("Start testing...")
    model.eval()
    meters = AverageMeterGroup()
    start_time = time.time()
    with torch.no_grad():
        for step, (inputs, targets) in enumerate(loader):
            inputs, targets = inputs.to('cuda'), targets.to('cuda')
            logits = model(inputs)
            loss = criterion(logits, targets)
            metrics = accuracy(logits, targets)
            metrics["loss"] = loss.item()
            meters.update(metrics)
            if step % log_freq == 0 or step + 1 == len(loader):
                logger.info("Valid Step [%d/%d] time %.3fs acc1 %.4f acc5 %.4f loss %.4f",
                            step + 1, len(loader), time.time() - start_time,
                            meters.acc1.avg, meters.acc5.avg, meters.loss.avg)
    return meters.acc1.avg


def evaluate_acc(class_cls, criterion, args, train_dataset, val_dataset):
    model = class_cls()
    with original_state_dict_hooks(model):
        model.load_state_dict(load_and_parse_state_dict(args.checkpoint), strict=False)
    model.cuda()
    
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.train_batch_size, num_workers=args.workers)
    test_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.test_batch_size, num_workers=args.workers)

    acc_before = test_acc(model, criterion, args.log_frequency, test_loader)
    nni.report_intermediate_result(acc_before)

    retrain_bn(model, criterion, args.train_iters, args.log_frequency, train_loader)
    acc = test_acc(model, criterion, args.log_frequency, test_loader)
    assert isinstance(acc, float)
    nni.report_intermediate_result(acc)
    nni.report_final_result(acc)


class LatencyFilter:
    def __init__(self, threshold, predictor, predictor_version=None, reverse=False):
        """
        Filter the models according to predicted latency. If the predicted latency of the ir model is larger than
        the given threshold, the ir model will be filtered and will not be considered as a searched architecture.

        Parameters
        ----------
        threshold: `float`
            the threshold of latency
        config, hardware:
            determine the targeted device
        reverse: `bool`
            if reverse is `False`, then the model returns `True` when `latency < threshold`,
            else otherwise
        """
        self.predictors = load_latency_predictor(predictor, predictor_version)
        self.threshold = threshold

    def __call__(self, ir_model):
        latency = self.predictors.predict(ir_model, 'nni-ir')
        return latency < self.threshold


def _main():
    parser = argparse.ArgumentParser("SPOS Evolutional Search")
    parser.add_argument("--port", type=int, default=8084)
    parser.add_argument("--imagenet-dir", type=str, default="./data/imagenet")
    parser.add_argument("--checkpoint", type=str, default="./data/checkpoint-150000.pth.tar")
    parser.add_argument("--spos-preprocessing", action="store_true", default=False,
                        help="When true, image values will range from 0 to 255 and use BGR "
                             "(as in original repo).")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--workers", type=int, default=6)
    parser.add_argument("--train-batch-size", type=int, default=128)
    parser.add_argument("--train-iters", type=int, default=200)
    parser.add_argument("--test-batch-size", type=int, default=512)
    parser.add_argument("--log-frequency", type=int, default=10)
    parser.add_argument("--label-smoothing", type=float, default=0.1)
    parser.add_argument("--evolution-sample-size", type=int, default=10)
    parser.add_argument("--evolution-population-size", type=int, default=50)
    parser.add_argument("--evolution-cycles", type=int, default=10)
    parser.add_argument("--latency-filter", type=str, default=None,
                        help="Apply latency filter by calling the name of the applied hardware.")
    parser.add_argument("--latency-threshold", type=float, default=100)

    args = parser.parse_args()

    # use a fixed set of image will improve the performance
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)
    torch.backends.cudnn.deterministic = True

    assert torch.cuda.is_available()

    base_model = ShuffleNetV2OneShot()
    criterion = CrossEntropyLabelSmooth(1000, args.label_smoothing)
    if args.spos_preprocessing:
        # ``nni.trace`` is used to make transforms serializable, so that the trials can run other processes or on remote servers.
        trans = nni.trace(transforms.Compose)([
            nni.trace(transforms.RandomResizedCrop)(224),
            nni.trace(transforms.ColorJitter)(brightness=0.4, contrast=0.4, saturation=0.4),
            nni.trace(transforms.RandomHorizontalFlip)(0.5),
            nni.trace(ToBGRTensor)(),
        ])
    else:
        # ``nni.trace`` is used to make transforms serializable, so that the trials can run other processes or on remote servers.
        trans = nni.trace(transforms.Compose)([
            nni.trace(transforms.RandomResizedCrop)(224),
            nni.trace(transforms.ToTensor)()
        ])
    train_dataset = nni.trace(datasets.ImageNet)(args.imagenet_dir, split='train', transform=trans)
    val_dataset = nni.trace(datasets.ImageNet)(args.imagenet_dir, split='val', transform=trans)

    if args.latency_filter:
        latency_filter = LatencyFilter(threshold=args.latency_threshold, predictor=args.latency_filter)
    else:
        latency_filter = None

    evaluator = FunctionalEvaluator(evaluate_acc, criterion=criterion, args=args, train_dataset=train_dataset, val_dataset=val_dataset)
    evolution_strategy = strategy.RegularizedEvolution(
        model_filter=latency_filter,
        sample_size=args.evolution_sample_size, population_size=args.evolution_population_size, cycles=args.evolution_cycles)
    exp = RetiariiExperiment(base_model, evaluator, strategy=evolution_strategy)

    exp_config = RetiariiExeConfig('local')
    exp_config.trial_concurrency = 2
    exp_config.trial_gpu_number = 1
    exp_config.max_trial_number = args.evolution_cycles
    exp_config.training_service.use_active_gpu = False
    exp_config.execution_engine = 'base'
    exp_config.dummy_input = [1, 3, 224, 224]

    exp.run(exp_config, args.port)

    print('Exported models:')
    for i, model in enumerate(exp.export_top_models(formatter='dict')):
        print(model)
        with open(f'architecture_final_{i}.json', 'w') as f: 
            json.dump(get_archchoice_by_model(model), f, indent=4)

if __name__ == "__main__":
    _main()