def generate_tasks(self, task_result: TaskResult) -> List[Task]:
# initial/update temp config list
if not self._temp_config_list:
self._init_temp_config_list()
else:
score = self._tasks[task_result.task_id].score
assert score is not None, 'SimulatedAnnealingTaskGenerator need each score is not None.'
if not self._current_sparsity_list:
self._current_sparsity_list = deepcopy(self._temp_sparsity_list)
self._current_score = score
else:
delta_E = np.abs(score - self._current_score)
probability = np.exp(-1 * delta_E / self.current_temperature)
if self._current_score < score or np.random.uniform(0, 1) < probability:
self._current_score = score
self._current_sparsity_list = deepcopy(self._temp_sparsity_list)
self.current_temperature *= self.cool_down_rate
if self.current_temperature < self.stop_temperature:
return []
self._update_with_perturbations()
task_id = self._task_id_candidate
new_config_list = self._recover_real_sparsity(deepcopy(self._temp_config_list))
config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))
with Path(config_list_path).open('w') as f:
json_tricks.dump(new_config_list, f, indent=4)
task = Task(task_id, self.temp_model_path, self.temp_masks_path, config_list_path)
self._tasks[task_id] = task
self._task_id_candidate += 1
return [task]
def generate_tasks(self, task_result: TaskResult) -> List[Task]:
compact_model = task_result.compact_model
compact_model_masks = task_result.compact_model_masks
# save intermediate result
model_path = Path(self._intermediate_result_dir,
'{}_compact_model.pth'.format(task_result.task_id))
masks_path = Path(
self._intermediate_result_dir,
'{}_compact_model_masks.pth'.format(task_result.task_id))
torch.save(compact_model, model_path)
torch.save(compact_model_masks, masks_path)
# get current2origin_sparsity and compact2origin_sparsity
origin_model = torch.load(self._origin_model_path)
current2origin_sparsity, compact2origin_sparsity, _ = compute_sparsity(
origin_model, compact_model, compact_model_masks,
self.target_sparsity)
_logger.debug(
'\nTask %s total real sparsity compared with original model is:\n%s',
str(task_result.task_id),
json_tricks.dumps(current2origin_sparsity, indent=4))
if task_result.task_id != 'origin':
self._tasks[task_result.task_id].state[
'current2origin_sparsity'] = current2origin_sparsity
# if reach the total_iteration, no more task will be generated
if self.current_iteration > self.total_iteration:
return []
task_id = self._task_id_candidate
new_config_list = self.generate_config_list(self.target_sparsity,
self.current_iteration,
compact2origin_sparsity)
new_config_list = self.allocate_sparsity(new_config_list,
compact_model,
compact_model_masks)
config_list_path = Path(self._intermediate_result_dir,
'{}_config_list.json'.format(task_id))
with Path(config_list_path).open('w') as f:
json_tricks.dump(new_config_list, f, indent=4)
task = Task(task_id, model_path, masks_path, config_list_path)
self._tasks[task_id] = task
self._task_id_candidate += 1
self.current_iteration += 1
return [task]
def pruning_one_step_normal(self, task: Task) -> TaskResult:
"""
generate masks -> speedup -> finetune -> evaluate
"""
model, masks, config_list = task.load_data()
self.pruner.reset(model, config_list)
self.pruner.load_masks(masks)
# pruning model
compact_model, pruner_generated_masks = self.pruner.compress()
compact_model_masks = deepcopy(pruner_generated_masks)
# show the pruning effect
self.pruner.show_pruned_weights()
self.pruner._unwrap_model()
# speedup
if self.speedup and task.speedup:
ModelSpeedup(compact_model, self.dummy_input,
pruner_generated_masks).speedup_model()
compact_model_masks = {}
# finetune
if self.finetuner is not None and task.finetune:
if self.speedup:
self.finetuner(compact_model)
else:
self.pruner._wrap_model()
self.finetuner(compact_model)
self.pruner._unwrap_model()
# evaluate
if self.evaluator is not None and task.evaluate:
if self.speedup:
score = self.evaluator(compact_model)
else:
self.pruner._wrap_model()
score = self.evaluator(compact_model)
self.pruner._unwrap_model()
else:
score = None
# clear model references
self.pruner.clear_model_references()
return TaskResult(task.task_id, compact_model, compact_model_masks,
pruner_generated_masks, score)
def generate_tasks(self, task_result: TaskResult) -> List[Task]:
# append experience & update agent policy
if self.action is not None:
action, reward, observation, done = self.env.step(self.action, task_result.compact_model)
self.T.append([reward, self.observation, observation, self.action, done])
self.observation = observation.copy()
if done:
assert task_result.score is not None, 'task_result.score should not be None if environment is done.'
final_reward = task_result.score - 1
# agent observe and update policy
for _, s_t, s_t1, a_t, d_t in self.T:
self.agent.observe(final_reward, s_t, s_t1, a_t, d_t)
if self.current_episode > self.warmup_episode:
self.agent.update_policy()
self.current_episode += 1
self.T = []
self.action = None
self.observation = None
# update current2origin_sparsity in log file
origin_model = torch.load(self._origin_model_path)
compact_model = task_result.compact_model
compact_model_masks = task_result.compact_model_masks
current2origin_sparsity, _, _ = compute_sparsity(origin_model, compact_model, compact_model_masks, self.temp_config_list)
self._tasks[task_result.task_id].state['current2origin_sparsity'] = current2origin_sparsity
current2origin_sparsity, _, _ = compute_sparsity(origin_model, compact_model, compact_model_masks, self.config_list_copy)
self._tasks[task_result.task_id].state['current_total_sparsity'] = current2origin_sparsity
flops, params, _ = count_flops_params(compact_model, self.dummy_input, verbose=False)
self._tasks[task_result.task_id].state['current_flops'] = '{:.2f} M'.format(flops / 1e6)
self._tasks[task_result.task_id].state['current_params'] = '{:.2f} M'.format(params / 1e6)
# generate new action
if self.current_episode < self.total_episode:
if self.observation is None:
self.observation = self.env.reset().copy()
self.temp_config_list = []
compact_model = torch.load(self._origin_model_path)
compact_model_masks = torch.load(self._origin_masks_path)
else:
compact_model = task_result.compact_model
compact_model_masks = task_result.compact_model_masks
if self.current_episode <= self.warmup_episode:
action = self.agent.random_action()
else:
action = self.agent.select_action(self.observation, episode=self.current_episode)
action = action.tolist()[0]
self.action = self.env.correct_action(action, compact_model)
sub_config_list = [{'op_names': [self.env.current_op_name], 'total_sparsity': self.action}]
self.temp_config_list.extend(sub_config_list)
task_id = self._task_id_candidate
if self.env.is_first_layer() or self.env.is_final_layer():
task_config_list = self.temp_config_list
else:
task_config_list = sub_config_list
config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))
with Path(config_list_path).open('w') as f:
json_tricks.dump(task_config_list, f, indent=4)
model_path = Path(self._intermediate_result_dir, '{}_compact_model.pth'.format(task_result.task_id))
masks_path = Path(self._intermediate_result_dir, '{}_compact_model_masks.pth'.format(task_result.task_id))
torch.save(compact_model, model_path)
torch.save(compact_model_masks, masks_path)
task = Task(task_id, model_path, masks_path, config_list_path)
if not self.env.is_final_layer():
task.finetune = False
task.evaluate = False
self._tasks[task_id] = task
self._task_id_candidate += 1
return [task]
else:
return []
