diff --git a/paddlex/cv/models/slim/prune.py b/paddlex/cv/models/slim/prune.py index 4ff3e237d13a156f96f21360a5cb8393dbdd9e40..f2dc8da10eb08880c3a9ed06a4b40e6b593c5bba 100644 --- a/paddlex/cv/models/slim/prune.py +++ b/paddlex/cv/models/slim/prune.py @@ -104,7 +104,7 @@ def sensitivity(program, return sensitivities -def channel_prune(program, prune_names, prune_ratios, place, only_graph=False): +def channel_prune(program, prune_names, prune_ratios, place, only_graph=False, scope=None): """通道裁剪。 Args: @@ -134,7 +134,8 @@ def channel_prune(program, prune_names, prune_ratios, place, only_graph=False): pruned_num = int(round(origin_num * (ratio))) prune_ratios[index] = ratio index += 1 - scope = fluid.global_scope() + if scope is None: + scope = fluid.global_scope() pruner = Pruner() program, _, _ = pruner.prune( program, @@ -175,12 +176,12 @@ def prune_program(model, prune_params_ratios=None): prune_params_ratios[prune_name] for prune_name in prune_names ] model.train_prog = channel_prune(train_prog, prune_names, prune_ratios, - place) + place, scope=model.scope) model.test_prog = channel_prune( - eval_prog, prune_names, prune_ratios, place, only_graph=True) + eval_prog, prune_names, prune_ratios, place, only_graph=True, scope=model.scope) -def update_program(program, model_dir, place): +def update_program(program, model_dir, place, scope=None): """根据裁剪信息更新Program和参数。 Args: @@ -197,10 +198,12 @@ def update_program(program, model_dir, place): shapes = yaml.load(f.read(), Loader=yaml.Loader) for param, shape in shapes.items(): graph.var(param).set_shape(shape) + if scope is None: + scope = fluid.global_scope() for block in program.blocks: for param in block.all_parameters(): if param.name in shapes: - param_tensor = fluid.global_scope().find_var( + param_tensor = scope.find_var( param.name).get_tensor() param_tensor.set( np.zeros(list(shapes[param.name])).astype('float32'), @@ -293,7 +296,7 @@ def get_params_ratios(sensitivities_file, eval_metric_loss=0.05): return params_ratios -def cal_model_size(program, place, sensitivities_file, eval_metric_loss=0.05): +def cal_model_size(program, place, sensitivities_file, eval_metric_loss=0.05, scope=None): """在可容忍的精度损失下,计算裁剪后模型大小相对于当前模型大小的比例。 Args: @@ -326,7 +329,8 @@ def cal_model_size(program, place, sensitivities_file, eval_metric_loss=0.05): list(prune_params_ratios.keys()), list(prune_params_ratios.values()), place, - only_graph=True) + only_graph=True, + scope=scope) origin_size = 0 new_size = 0 for var in program.list_vars(): diff --git a/paddlex/cv/models/slim/prune_config.py b/paddlex/cv/models/slim/prune_config.py index e9019f21e874e3960fd19b133c89dc607ac1c147..6bfe35375f87dfbdbe32b2bdeb6c0bc6cff86706 100644 --- a/paddlex/cv/models/slim/prune_config.py +++ b/paddlex/cv/models/slim/prune_config.py @@ -171,10 +171,14 @@ def get_prune_params(model): model_type.startswith('ShuffleNetV2'): for block in program.blocks: for param in block.all_parameters(): - pd_var = fluid.global_scope().find_var(param.name) - pd_param = pd_var.get_tensor() - if len(np.array(pd_param).shape) == 4: - prune_names.append(param.name) + pd_var = model.scope.find_var(param.name) + try: + pd_param = pd_var.get_tensor() + if len(np.array(pd_param).shape) == 4: + prune_names.append(param.name) + except Exception as e: + print("None Tensor Name: ", param.name) + print("Error message: {}".format(e)) if model_type == 'AlexNet': prune_names.remove('conv5_weights') if model_type == 'ShuffleNetV2': diff --git a/paddlex/cv/models/slim/visualize.py b/paddlex/cv/models/slim/visualize.py index 4be6721632cd7c8d26309cedb686466d2c0ec776..a8cb444d840fcdcb9b5dda8f77ab73d20a96b03a 100644 --- a/paddlex/cv/models/slim/visualize.py +++ b/paddlex/cv/models/slim/visualize.py @@ -42,7 +42,7 @@ def visualize(model, sensitivities_file, save_dir='./'): y = list() for loss_thresh in tqdm.tqdm(list(np.arange(0.05, 1, 0.05))): prune_ratio = 1 - cal_model_size( - program, place, sensitivities_file, eval_metric_loss=loss_thresh) + program, place, sensitivities_file, eval_metric_loss=loss_thresh, scope=model.scope) x.append(prune_ratio) y.append(loss_thresh) plt.plot(x, y, color='green', linewidth=0.5, marker='o', markersize=3)