sanas_darts_space.py

import os
import sys
sys.path.append('..')
import numpy as np
import argparse
import ast
import time
import argparse
import ast
import logging
import paddle.fluid as fluid
from paddleslim.nas import SANAS
from paddleslim.common import get_logger
import darts_cifar10_reader as reader

_logger = get_logger(__name__, level=logging.INFO)

auxiliary = True
auxiliary_weight = 0.4
trainset_num = 50000
lr = 0.025
momentum = 0.9
weight_decay = 0.0003
drop_path_probility = 0.2


class AvgrageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.avg = 0
        self.sum = 0
        self.cnt = 0

    def update(self, val, n=1):
        self.sum += val * n
        self.cnt += n
        self.avg = self.sum / self.cnt


def count_parameters_in_MB(all_params, prefix='model'):
    parameters_number = 0
    for param in all_params:
        if param.name.startswith(
                prefix) and param.trainable and 'aux' not in param.name:
            parameters_number += np.prod(param.shape)
    return parameters_number / 1e6


def create_data_loader(image_shape, is_train, args):
    image = fluid.data(
        name="image", shape=[None] + image_shape, dtype="float32")
    label = fluid.data(name="label", shape=[None, 1], dtype="int64")
    data_loader = fluid.io.DataLoader.from_generator(
        feed_list=[image, label],
        capacity=64,
        use_double_buffer=True,
        iterable=True)
    drop_path_prob = ''
    drop_path_mask = ''
    if is_train:
        drop_path_prob = fluid.data(
            name="drop_path_prob", shape=[args.batch_size, 1], dtype="float32")
        drop_path_mask = fluid.data(
            name="drop_path_mask",
            shape=[args.batch_size, 20, 4, 2],
            dtype="float32")

    return data_loader, image, label, drop_path_prob, drop_path_mask


def build_program(main_program, startup_program, image_shape, archs, args,
                  is_train):
    with fluid.program_guard(main_program, startup_program):
        data_loader, data, label, drop_path_prob, drop_path_mask = create_data_loader(
            image_shape, is_train, args)
        logits, logits_aux = archs(data, drop_path_prob, drop_path_mask,
                                   is_train, 10)
        top1 = fluid.layers.accuracy(input=logits, label=label, k=1)
        top5 = fluid.layers.accuracy(input=logits, label=label, k=5)
        loss = fluid.layers.reduce_mean(
            fluid.layers.softmax_with_cross_entropy(logits, label))

        if is_train:
            if auxiliary:
                loss_aux = fluid.layers.reduce_mean(
                    fluid.layers.softmax_with_cross_entropy(logits_aux, label))
                loss = loss + auxiliary_weight * loss_aux
            step_per_epoch = int(trainset_num / args.batch_size)
            learning_rate = fluid.layers.cosine_decay(lr, step_per_epoch,
                                                      args.retain_epoch)
            fluid.clip.set_gradient_clip(
                clip=fluid.clip.GradientClipByGlobalNorm(clip_norm=5.0))
            optimizer = fluid.optimizer.MomentumOptimizer(
                learning_rate,
                momentum,
                regularization=fluid.regularizer.L2DecayRegularizer(
                    weight_decay))
            optimizer.minimize(loss)
            outs = [loss, top1, top5, learning_rate]
        else:
            outs = [loss, top1, top5]
    return outs, data_loader


def train(main_prog, exe, epoch_id, train_loader, fetch_list, args):
    loss = AvgrageMeter()
    top1 = AvgrageMeter()
    top5 = AvgrageMeter()
    for step_id, data in enumerate(train_loader()):
        devices_num = len(data)
        if drop_path_probility > 0:
            feed = []
            for device_id in range(devices_num):
                image = data[device_id]['image']
                label = data[device_id]['label']
                drop_path_prob = np.array(
                    [[drop_path_probility * epoch_id / args.retain_epoch]
                     for i in range(args.batch_size)]).astype(np.float32)
                drop_path_mask = 1 - np.random.binomial(
                    1, drop_path_prob[0],
                    size=[args.batch_size, 20, 4, 2]).astype(np.float32)
                feed.append({
                    "image": image,
                    "label": label,
                    "drop_path_prob": drop_path_prob,
                    "drop_path_mask": drop_path_mask
                })
        else:
            feed = data
        loss_v, top1_v, top5_v, lr = exe.run(
            main_prog, feed=feed, fetch_list=[v.name for v in fetch_list])
        loss.update(loss_v, args.batch_size)
        top1.update(top1_v, args.batch_size)
        top5.update(top5_v, args.batch_size)
        if step_id % 10 == 0:
            _logger.info(
                "Train Epoch {}, Step {}, Lr {:.8f}, loss {:.6f}, acc_1 {:.6f}, acc_5 {:.6f}".
                format(epoch_id, step_id, lr[0], loss.avg[0], top1.avg[0],
                       top5.avg[0]))
    return top1.avg[0]


def valid(main_prog, exe, epoch_id, valid_loader, fetch_list, args):
    loss = AvgrageMeter()
    top1 = AvgrageMeter()
    top5 = AvgrageMeter()
    for step_id, data in enumerate(valid_loader()):
        loss_v, top1_v, top5_v = exe.run(
            main_prog, feed=data, fetch_list=[v.name for v in fetch_list])
        loss.update(loss_v, args.batch_size)
        top1.update(top1_v, args.batch_size)
        top5.update(top5_v, args.batch_size)
        if step_id % 10 == 0:
            _logger.info(
                "Valid Epoch {}, Step {}, loss {:.6f}, acc_1 {:.6f}, acc_5 {:.6f}".
                format(epoch_id, step_id, loss.avg[0], top1.avg[0], top5.avg[
                    0]))
    return top1.avg[0]


def search(config, args, image_size, is_server=True):
    if is_server:
        ### start a server and a client
        sa_nas = SANAS(
            config,
            server_addr=(args.server_address, args.port),
            search_steps=args.search_steps,
            is_server=True)
    else:
        ### start a client
        sa_nas = SANAS(
            config,
            server_addr=(args.server_address, args.port),
            init_temperature=init_temperature,
            is_server=False)

    image_shape = [3, image_size, image_size]
    for step in range(args.search_steps):
        archs = sa_nas.next_archs()[0]

        train_program = fluid.Program()
        test_program = fluid.Program()
        startup_program = fluid.Program()
        train_fetch_list, train_loader = build_program(
            train_program,
            startup_program,
            image_shape,
            archs,
            args,
            is_train=True)

        current_params = count_parameters_in_MB(
            train_program.global_block().all_parameters(), 'cifar10')
        _logger.info('step: {}, current_params: {}M'.format(step,
                                                            current_params))
        if current_params > float(3.77):
            continue

        test_fetch_list, test_loader = build_program(
            test_program,
            startup_program,
            image_shape,
            archs,
            args,
            is_train=False)
        test_program = test_program.clone(for_test=True)

        place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
        exe = fluid.Executor(place)
        exe.run(startup_program)

        train_reader = reader.train_valid(
            batch_size=args.batch_size, is_train=True, is_shuffle=True)
        test_reader = reader.train_valid(
            batch_size=args.batch_size, is_train=False, is_shuffle=False)

        train_loader.set_batch_generator(train_reader, places=place)
        test_loader.set_batch_generator(test_reader, places=place)

        build_strategy = fluid.BuildStrategy()
        train_compiled_program = fluid.CompiledProgram(
            train_program).with_data_parallel(
                loss_name=train_fetch_list[0].name,
                build_strategy=build_strategy)

        valid_top1_list = []
        for epoch_id in range(args.retain_epoch):
            train_top1 = train(train_compiled_program, exe, epoch_id,
                               train_loader, train_fetch_list, args)
            _logger.info("TRAIN: step: {}, Epoch {}, train_acc {:.6f}".format(
                step, epoch_id, train_top1))
            valid_top1 = valid(test_program, exe, epoch_id, test_loader,
                               test_fetch_list, args)
            _logger.info("TEST: Epoch {}, valid_acc {:.6f}".format(epoch_id,
                                                                   valid_top1))
            valid_top1_list.append(valid_top1)
        sa_nas.reward(float(valid_top1_list[-1] + valid_top1_list[-2]) / 2)


def final_test(config, args, image_size, token=None):
    assert token != None, "If you want to start a final experiment, you must input a token."
    sa_nas = SANAS(
        config, server_addr=(args.server_address, args.port), is_server=True)

    image_shape = [3, image_size, image_size]
    archs = sa_nas.tokens2arch(token)[0]

    train_program = fluid.Program()
    test_program = fluid.Program()
    startup_program = fluid.Program()
    train_fetch_list, train_loader = build_program(
        train_program,
        startup_program,
        image_shape,
        archs,
        args,
        is_train=True)

    current_params = count_parameters_in_MB(
        train_program.global_block().all_parameters(), 'cifar10')
    _logger.info('current_params: {}M'.format(current_params))
    test_fetch_list, test_loader = build_program(
        test_program,
        startup_program,
        image_shape,
        archs,
        args,
        is_train=False)
    test_program = test_program.clone(for_test=True)

    place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
    exe = fluid.Executor(place)
    exe.run(startup_program)

    train_reader = reader.train_valid(
        batch_size=args.batch_size, is_train=True, is_shuffle=True, args=args)
    test_reader = reader.train_valid(
        batch_size=args.batch_size,
        is_train=False,
        is_shuffle=False,
        args=args)

    train_loader.set_batch_generator(train_reader, places=place)
    test_loader.set_batch_generator(test_reader, places=place)

    build_strategy = fluid.BuildStrategy()
    train_compiled_program = fluid.CompiledProgram(
        train_program).with_data_parallel(
            loss_name=train_fetch_list[0].name, build_strategy=build_strategy)

    valid_top1_list = []
    for epoch_id in range(args.retain_epoch):
        train_top1 = train(train_compiled_program, exe, epoch_id, train_loader,
                           train_fetch_list, args)
        _logger.info("TRAIN: Epoch {}, train_acc {:.6f}".format(epoch_id,
                                                                train_top1))
        valid_top1 = valid(test_program, exe, epoch_id, test_loader,
                           test_fetch_list, args)
        _logger.info("TEST: Epoch {}, valid_acc {:.6f}".format(epoch_id,
                                                               valid_top1))
        valid_top1_list.append(valid_top1)

        output_dir = os.path.join('darts_output', str(epoch_id))
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        fluid.io.save_persistables(exe, output_dir, main_program=train_program)


if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description='SA NAS MobileNetV2 cifar10 argparase')
    parser.add_argument(
        '--use_gpu',
        type=ast.literal_eval,
        default=True,
        help='Whether to use GPU in train/test model.')
    parser.add_argument(
        '--batch_size', type=int, default=96, help='batch size.')
    parser.add_argument(
        '--is_server',
        type=ast.literal_eval,
        default=True,
        help='Whether to start a server.')
    parser.add_argument(
        '--server_address', type=str, default="", help='server ip.')
    parser.add_argument('--port', type=int, default=8881, help='server port')
    parser.add_argument(
        '--retain_epoch', type=int, default=30, help='epoch for each token.')
    parser.add_argument('--token', type=int, nargs='+', help='final token.')
    parser.add_argument(
        '--search_steps',
        type=int,
        default=200,
        help='controller server number.')
    args = parser.parse_args()
    print(args)

    image_size = 32

    config = [('DartsSpace')]

    if args.token == None:
        search(config, args, image_size, is_server=args.is_server)
    else:
        final_test(config, args, image_size, token=args.token)