Dubhe/dubhe-tadl/spos/tester.py

import argparse
# import logging
import random
import time
from itertools import cycle
import json


import sys
sys.path.insert(0,"/home/hanjiayi/Document/nni")
sys.path.insert(0,"/home/hanjiayi/Document/nni/examples")
sys.path.insert(0,"/home/hanjiayi/Document/nni/nni/algorithms")
import os

# 指定写入文件的环境变量，原实现于nnictl_utils.py的search_space_auto_gen
os.environ["NNI_GEN_SEARCH_SPACE"] = "auto_gen_search_space.json"

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"


import numpy as np
import torch
import torch.nn as nn
from nni.algorithms.nas.pytorch.classic_nas import get_and_apply_next_architecture
from nni.nas.pytorch.fixed import apply_fixed_architecture
from nni.nas.pytorch.utils import AverageMeterGroup

from dataloader import get_imagenet_iter_dali
from network import ShuffleNetV2OneShot, load_and_parse_state_dict
from utils import CrossEntropyLabelSmooth, accuracy


# logger = logging.getLogger("nni.spos.tester")     # "nni.spos.tester"
print("Evolution Beginning...")

def retrain_bn(model, criterion, max_iters, log_freq, loader):
    with torch.no_grad():
        # logger.info("Clear BN statistics...")
        print("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)...")
        print("Train BN with training set (BN sanitize)...")
        model.train()
        meters = AverageMeterGroup()
        start_time = time.time()

        for step in range(max_iters):
            inputs, targets = next(loader)
            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 time %.3fs ", step + 1, max_iters, meters, time.time() - start_time)
                print("Train Step [%d/%d] %s time %.3fs "% (step + 1, max_iters, meters, time.time() - start_time))

def test_acc(model, criterion, log_freq, loader):
    # logger.info("Start testing...")
    print("start testing...")
    model.eval()
    meters = AverageMeterGroup()
    start_time = time.time()

    with torch.no_grad():
        for step, (inputs, targets) in enumerate(loader):
            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)
                print("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))
            if step>len(loader):     # 遍历一遍就停止
                break
    return meters.acc1.avg


def evaluate_acc(model, criterion, args, loader_train, loader_test):

    retrain_bn(model, criterion, args.train_iters, args.log_frequency, loader_train)     # todo
    acc = test_acc(model, criterion, args.log_frequency, loader_test)
    assert isinstance(acc, float)
    torch.cuda.empty_cache()
    return acc


if __name__ == "__main__":
    parser = argparse.ArgumentParser("SPOS Candidate Tester")
    parser.add_argument("--imagenet-dir", type=str, default="/mnt/local/hanjiayi/imagenet")     # ./data/imagenet
    parser.add_argument("--checkpoint", type=str, default="./data/checkpoint-150000.pth.tar")     # ./data/checkpoint-150000.pth.tar
    parser.add_argument("--spos-preprocessing", default=True,
                        help="When true, image values will range from 0 to 255 and use BGR "
                             "(as in original repo).")     # , action="store_true"
    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)    # nni中为512，官方repo为200
    parser.add_argument("--log-frequency", type=int, default=10)
    parser.add_argument("--architecture", type=str, default="./architecture_final.json", help="load the file to retrain or eval")
    parser.add_argument("--mode", type=str, default="gen", help="there are two modes here: gen mode for generating architecture, and evl mode for evaluation model")

    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()

    model = ShuffleNetV2OneShot()
    criterion = CrossEntropyLabelSmooth(1000, 0.1)


    if args.mode == "gen":
        get_and_apply_next_architecture(model)
        model.load_state_dict(load_and_parse_state_dict(filepath=args.checkpoint))

    else:     # evaluate the model
        print("## test&retrain -- load model ## begin to load model")
        model.load_state_dict(load_and_parse_state_dict(filepath=args.checkpoint))
        print("## test&retrain -- load model ## model loaded")

        print("## test&retrain -- apply architecture ## begin to apply architecture to model")
        apply_fixed_architecture(model, args.architecture)
        print("## test&retrain -- apply architecture ## architecture applied")

    model.cuda(0)
    print("## test&retrain -- load train data ## begin to load train data")
    train_loader = get_imagenet_iter_dali("train", args.imagenet_dir, args.train_batch_size, args.workers,
                                          spos_preprocessing=args.spos_preprocessing,
                                          seed=args.seed, device_id=0)
    print("## test&retrain -- load train data ## train data loaded")

    print("## test&retrain -- load test data ## begin to load test data")
    val_loader = get_imagenet_iter_dali("val", args.imagenet_dir, args.test_batch_size, args.workers,
                                        spos_preprocessing=args.spos_preprocessing, shuffle=True,
                                        seed=args.seed, device_id=0)
    print("## test&retrain -- load test date ## test data loaded")

    train_loader = cycle(train_loader)
    acc = evaluate_acc(model, criterion, args, train_loader, val_loader)

    # 把模型最终的准确率写入一个文件中
    os.makedirs("./acc", exist_ok=True)
    with open("./acc/{}".format(args.architecture[-12:]), "w") as f:     # [-12:] 代表没有路径的文件名
        # {filename1: acc,
        #  filename2: acc,
        #  000_000.json: acc,
        #  000_001.json: acc,
        #  ......
        #  }
        json.dump({args.architecture: acc}, f)