import os import re import sys import json import pickle import logging import subprocess import numpy as np from collections import deque from evaluator import Evaluator from network import ShuffleNetV2OneShot, PARSED_FLOPS LAYER_CHOICE = "layer_choice" INPUT_CHOICE = "input_choice" _logger = logging.getLogger(__name__) class SPOSEvolution: """ SPOS evolution tuner. Parameters ---------- max_epochs : int Maximum number of epochs to run. num_select : int Number of survival candidates of each epoch. num_population : int Number of candidates at the start of each epoch. If candidates generated by crossover and mutation are not enough, the rest will be filled with random candidates. m_prob : float The probability of mutation. num_crossover : int Number of candidates generated by crossover in each epoch. num_mutation : int Number of candidates generated by mutation in each epoch. """ def __init__(self, max_epochs=20, num_select=10, num_population=50, m_prob=0.1, num_crossover=25, num_mutation=25, epoch=0): assert num_population >= num_select self.max_epochs = max_epochs self.num_select = num_select self.num_population = num_population self.m_prob = m_prob self.num_crossover = num_crossover self.num_mutation = num_mutation self.epoch = epoch self.search_space = None self.random_state = np.random.RandomState(0) # self.evl = Evaluator() # async status self._to_evaluate_queue = deque() self._sending_parameter_queue = deque() self._pending_result_ids = set() self._reward_dict = dict() self._id2candidate = dict() self._st_callback = None self.cand_path = "./checkpoints" self.acc_path = "./acc" self.candidates = [] if epoch == 0 else self.load_candidates() # 第一轮初始尚未有生成的种群 def load_candidates(self): # 从self.export_result()写入文件的候选模型，需要读入 # {"LayerChoice1": [false, false, false, true], ... } -> {"LayerChoice1": {"_idx":3, "_value":"3"}, ... } print("## evolution -- load ## begin to load candidates in evolution...\n") file_dir, _, files = next(os.walk(self.cand_path)) files = [i for i in files if "%03d_"%(self.epoch-1) in i] def get_true_index(l): return [i for i in range(len(l)) if l[i]][0] candidates = [] for file in files: with open(os.path.join(file_dir, file), "r") as f: candidate = json.load(f) # 转换成合适的形式 cand = {} for key, value in candidate.items(): v = get_true_index(value) value = {"_value":str(v), "_idx":int(v)} cand.update({key:value}) candidates.append(cand) print("## evolution -- load ## candidates loaded \n") return candidates def load_id2candidate(self): with open("./id2cand/%03d_id2candidate.json"%(self.epoch - 1), "r") as f: self.id2candidate = json.load(f) def update_search_space(self, search_space): """ Handle the initialization/update event of search space. """ print("## evolution -- update ## updating search space") self._search_space = search_space self._next_round() print("## evolution -- update ## search space updated") def _next_round(self): _logger.info("Epoch %d, generating...", self.epoch) if self.epoch == 0: self._get_random_population() self.export_results(self.candidates) self.evaluate_cands() # 评估全部的模型 else: self.load_id2candidate() self.receive_trial_result() best_candidates = self._select_top_candidates() if self.epoch >= self.max_epochs: return self.candidates = self._get_mutation(best_candidates) + self._get_crossover(best_candidates) self._get_random_population() self.export_results(self.candidates) self.evaluate_cands() # 评估全部的模型 self.epoch += 1 def _random_candidate(self): chosen_arch = dict() for key, val in self._search_space.items(): if val["_type"] == LAYER_CHOICE: choices = val["_value"] index = self.random_state.randint(len(choices)) chosen_arch[key] = {"_value": choices[index], "_idx": index} elif val["_type"] == INPUT_CHOICE: raise NotImplementedError("Input choice is not implemented yet.") return chosen_arch def _add_to_evaluate_queue(self, cand): _logger.info("Generate candidate %s, adding to eval queue.", self._get_architecture_repr(cand)) self._reward_dict[self._hashcode(cand)] = 0. self._to_evaluate_queue.append(cand) def _get_random_population(self): while len(self.candidates) < self.num_population: cand = self._random_candidate() if self._is_legal(cand): _logger.info("Random candidate generated.") self._add_to_evaluate_queue(cand) self.candidates.append(cand) def _get_crossover(self, best): result = [] for _ in range(10 * self.num_crossover): cand_p1 = best[self.random_state.randint(len(best))] cand_p2 = best[self.random_state.randint(len(best))] assert cand_p1.keys() == cand_p2.keys() cand = {k: cand_p1[k] if self.random_state.randint(2) == 0 else cand_p2[k] for k in cand_p1.keys()} if self._is_legal(cand): result.append(cand) self._add_to_evaluate_queue(cand) if len(result) >= self.num_crossover: break _logger.info("Found %d architectures with crossover.", len(result)) return result def _get_mutation(self, best): result = [] for _ in range(10 * self.num_mutation): cand = best[self.random_state.randint(len(best))].copy() mutation_sample = np.random.random_sample(len(cand)) for s, k in zip(mutation_sample, cand): if s < self.m_prob: choices = self._search_space[k]["_value"] index = self.random_state.randint(len(choices)) cand[k] = {"_value": choices[index], "_idx": index} if self._is_legal(cand): result.append(cand) self._add_to_evaluate_queue(cand) if len(result) >= self.num_mutation: break _logger.info("Found %d architectures with mutation.", len(result)) return result def _get_architecture_repr(self, cand): return re.sub(r"\".*?\": \{\"_idx\": (\d+), \"_value\": \".*?\"\}", r"\1", self._hashcode(cand)) def _is_legal(self, cand): if self._hashcode(cand) in self._reward_dict: return False return True # 将模型输出，并重训练、评估 def evaluate_cands(self): """ 1、对输出的模型进行重训练 2、对重训练后的模型进行评估 以上内容通过tester.py脚本完成 """ print("## evolution -- evaluate ## begin to evaluate candidates...") file_dir, _, files = next(os.walk(self.cand_path)) # 获取文件夹下的文件 files = [i for i in files if "%03d_"%self.epoch in i] for file in files: file = os.path.join(file_dir, file) # self.evl.eval_model(epoch=self.epoch, architecture=file) python_interpreter_path = sys.executable subprocess.run([python_interpreter_path,\ "evaluator.py", "--architecture", file, "--epoch", str(self.epoch)]) print("## evolution -- evaluate ## candidates evaluated") def _select_top_candidates(self): print("## evolution -- select ## begin to select top candidates...") reward_query = lambda cand: self._reward_dict[self._hashcode(cand)] _logger.info("All candidate rewards: %s", list(map(reward_query, self.candidates))) result = sorted(self.candidates, key=reward_query, reverse=True)[:self.num_select] _logger.info("Best candidate rewards: %s", list(map(reward_query, result))) print("## evolution -- select ## selected done") return result @staticmethod def _hashcode(d): return json.dumps(d, sort_keys=True) def _bind_and_send_parameters(self): """ There are two types of resources: parameter ids and candidates. This function is called at necessary times to bind these resources to send new trials with st_callback. """ result = [] while self._sending_parameter_queue and self._to_evaluate_queue: parameter_id = self._sending_parameter_queue.popleft() parameters = self._to_evaluate_queue.popleft() self._id2candidate[parameter_id] = parameters result.append(parameters) self._pending_result_ids.add(parameter_id) self._st_callback(parameter_id, parameters) _logger.info("Send parameter [%d] %s.", parameter_id, self._get_architecture_repr(parameters)) return result def generate_multiple_parameters(self, parameter_id_list, **kwargs): """ Callback function necessary to implement a tuner. This will put more parameter ids into the parameter id queue. """ if "st_callback" in kwargs and self._st_callback is None: self._st_callback = kwargs["st_callback"] for parameter_id in parameter_id_list: self._sending_parameter_queue.append(parameter_id) self._bind_and_send_parameters() return [] # always not use this. might induce problem of over-sending # def receive_trial_result(self, parameter_id, parameters, value, **kwargs): # """ # Callback function. Receive a trial result. # """ # _logger.info("Candidate %d, reported reward %f", parameter_id, value) # self._reward_dict[self._hashcode(self._id2candidate[parameter_id])] = value def receive_trial_result(self): # 获取并更新self._reward_dict file_dir, _, files = next(os.walk(self.acc_path)) files = [i for i in files if "%03d_"%(self.epoch-1) in i] # self.epoch-1: 读取上一轮的结果 acc_dict = {} for file in files: with open(os.path.join(file_dir, file), "r") as f: acc_dict.update(json.load(f)) # {"000_001.json":0.56} for key, value in acc_dict.items(): key = key.lstrip("./checkpoints/") # 删掉路径，仅保留文件名 self._reward_dict.update({self.id2candidate[key]: value}) # todo {self.id2candidate[key]: key} def trial_end(self, parameter_id, success, **kwargs): """ Callback function when a trial is ended and resource is released. """ self._pending_result_ids.remove(parameter_id) if not self._pending_result_ids and not self._to_evaluate_queue: # a new epoch now self._next_round() assert self._st_callback is not None self._bind_and_send_parameters() def export_results(self, result): """ Export a number of candidates to `checkpoints` dir. Parameters ---------- result : dict Chosen architectures to be exported. """ os.makedirs("checkpoints", exist_ok=True) os.makedirs("id2cand", exist_ok=True) self.id2candidate = {} for i, cand in enumerate(result): converted = dict() for cand_key, cand_val in cand.items(): onehot = [k == cand_val["_idx"] for k in range(len(self._search_space[cand_key]["_value"]))] converted[cand_key] = onehot with open(os.path.join("checkpoints", "%03d_%03d.json" % (self.epoch, i)), "w") as fp: json.dump(converted, fp) """ self.id2candidate: { 000_000.json: {"LayerChoice1": {"_values":3, "_idx":3}, "LayerChoice2": {"_values":2, "_idx":2}, ...} ...... } """ self.id2candidate.update({"%03d_%03d.json" % (self.epoch, i): json.dumps(result[i], sort_keys=True)}) with open("./id2cand/%03d_id2candidate.json"%self.epoch, "w") as f: json.dump(self.id2candidate, f) class EvolutionWithFlops(SPOSEvolution): """ This tuner extends the function of evolution tuner, by limiting the flops generated by tuner. Needs a function to examine the flops. """ def __init__(self, flops_limit=330E6, **kwargs): super().__init__(**kwargs) # self.model = ShuffleNetV2OneShot() self.flops_limit = flops_limit with open(os.path.join(os.path.dirname(__file__), "./data/op_flops_dict.pkl"), "rb") as fp: self._op_flops_dict = pickle.load(fp) def _is_legal(self, cand): if not super()._is_legal(cand): return False if self.get_candidate_flops(cand) > self.flops_limit: return False return True def get_candidate_flops(self, candidate): """ this method is the same with ShuffleNetV2OneShot.get_candidate_flops, but we dont need to initialize that class. """ conv1_flops = self._op_flops_dict["conv1"][(3, 16, 224, 224, 2)] rest_flops = self._op_flops_dict["rest_operation"][(640, 1000, 7, 7, 1)] total_flops = conv1_flops + rest_flops for k, m in candidate.items(): parsed_flops_dict = PARSED_FLOPS[k] if isinstance(m, dict): # to be compatible with classical nas format total_flops += parsed_flops_dict[m["_idx"]] else: total_flops += parsed_flops_dict[torch.max(m, 0)[1]] return total_flops