mindspore 2.3.0__cp39-cp39-win_amd64.whl → 2.4.1__cp39-cp39-win_amd64.whl

mindspore/safeguard/rewrite_obfuscation.py

@@ -15,13 +15,14 @@
 """obfuscate network based on rewrite interfaces."""
 import os
 import re
-import secrets
 from pathlib import Path
+from string import Template
+import numpy as np
 
+import mindspore as ms
 from mindspore import ops, nn
-from mindspore.common.tensor import Tensor
-from mindspore import log as logger
-from mindspore import load_checkpoint, save_checkpoint
+from mindspore import load_checkpoint, save_checkpoint, log
+from mindspore.ops import functional as F
 from mindspore.rewrite import SymbolTree, Node, NodeType, ScopedValue
 from mindspore.rewrite.parsers import ClassDefParser
 from mindspore.rewrite.parsers import ModuleParser
@@ -29,36 +30,127 @@ from mindspore.rewrite.parsers import ModuleParser
 OBF_RATIOS_LENGTH = 1
 MAX_OBF_RATIOS_NUM = 50
 OBF_RATIOS_WIDTH = 0
-OBF_RATIOS_INSERT_INDEX = 0
+
+_supported_ops = {
+    'mul': ops.Mul,
+    'matmul': ops.MatMul,
+    'invert': ops.Inv
+}
+
+_supported_config_type = [
+    'obf_metadata_config',
+    'weight_obf_config',
+    'network_obf_config'
+]
+
+_supported_metadata_type = [
+    'random',
+    'rearrange'
+]
+
+obf_medatadata_template = {
+    'name': 'obf_metadata',
+    'shape': [1,],
+    'type': 'random',
+    'save_metadata': True,
+    'metadata_op': 'invert'
+}
+
+weight_obf_template = {
+    'target': '',
+    'weight_obf_ops': [{'name': 'mul', 'input_x': 'weight', 'input_y': 'obf_metadata'}]
+}
+
+network_obf_template = {
+    'module': '',
+    'target': '',
+    'insert_new_input': [{'name': 'obf_metadata'}],
+    'insert_ops': [{'name': 'mul', 'input_x': 'weight', 'input_y': 'obf_metadata'}]
+}
+
+
+def _transform_target_modules(target_modules):
+    """transform target_modules to obf config"""
+    obf_config = {}
+    path = target_modules[0]
+    target_list = target_modules[1].split('|')
+    max_layers = 12
+    layers = []
+    obf_medatadata = obf_medatadata_template.copy()
+    if len(target_modules) >= 3:
+        obfuscate_layers = target_modules[2].split(':')
+        if obfuscate_layers[1] != 'all':
+            max_layers = int(obfuscate_layers[1])
+        layers = [i for i in range(0, max_layers)]
+        path_new = path.replace("blocks", "blocks/${layer}")
+        network_obf_template['insert_ops'][0]['input_y'] = "obf_metadata_${layer}"
+        weight_obf_template['weight_obf_ops'][0]['input_y'] = "obf_metadata_${layer}"
+        weight_obf_template['name'] = "obf_metadata_${layer}"
+        obf_medatadata['layers'] = layers
+    else:
+        path_new = path
+    obf_config['obf_metadata_config'] = []
+    obf_config['weight_obf_config'] = []
+    obf_config['network_obf_config'] = []
+    obf_config['obf_metadata_config'].append(obf_medatadata)
+
+    for name in target_list:
+        target_weight = path_new + '/' + name + '/weight'
+        target_bias = path_new + '/' + name + '/bias'
+        weight_obf = weight_obf_template.copy()
+        weight_obf['target'] = target_weight
+        bias_obf = weight_obf_template.copy()
+        bias_obf['target'] = target_bias
+        network_obf = network_obf_template.copy()
+        network_obf['module'] = '/' + path_new
+        network_obf['target'] = name
+        if not layers:
+            weight_obf['layers'] = layers
+            bias_obf['layers'] = layers
+            network_obf['layers'] = layers
+        obf_config['weight_obf_config'].append(weight_obf)
+        obf_config['weight_obf_config'].append(bias_obf)
+        obf_config['network_obf_config'].append(network_obf)
+    return obf_config
+
+
+def _get_op(op_name):
+    if op_name is None:
+        return None
+    if op_name not in _supported_ops:
+        raise KeyError(f"'op name' must be in {list(_supported_ops.keys())}, but got {op_name}.")
+    return _supported_ops[op_name]()
 
 
-def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./', obfuscate_scale=100):
+def obfuscate_ckpt(network, ckpt_files, target_modules=None, obf_config=None, saved_path='./', obfuscate_scale=100):
     """
-    obfuscate the plaintext checkpoint files. Usually used in conjunction with
-    :func:`mindspore.load_obf_params_into_net`.
-    interface.
+    Obfuscate the plaintext checkpoint files according to the obfuscation config.
 
     Args:
         network (nn.Cell): The original network that need to be obfuscated.
         ckpt_files (str): The directory path of original ckpt files.
-        target_modules (list[str]): The target module of network that need to be obfuscated. The first string
-            represents the network path of target module in original network, which should be in form of ``'A/B/C'``.
-            The second string represents the obfuscation target module, which should be in form of ``'D|E|F'``. For
-            example, thr target_modules of GPT2 can be ``['backbone/blocks/attention', 'dense1|dense2|dense3']``.
-            If target_modules has the third value, it should be in the format of 'obfuscate_layers:all' or
-            'obfuscate_layers:int', which represents the number of layers need to be obfuscated of duplicate layers
-            (such as transformer layers or resnet blocks). If target_modules is ``None``, the function would search
-            target modules by itself. If found, the searched target module would be used, otherwise suggested target
-            modules would be given with warning log. Default: ``None``.
+        target_modules (list[str]): The target ops that need to be obfuscated in the network. The first string
+            represents the network path of the target ops in the original network, which should be in form of
+            ``"A/B/C"``. The second string represents the names of multiple target ops in the same path, which
+            should be in form of ``"D|E|F"``. For example, the target_modules of GPT2 can be ``['backbone/blocks
+            /attention', 'dense1|dense2|dense3']``. If target_modules has the third value, it should be in the
+            format of 'obfuscate_layers:all' or 'obfuscate_layers:int', which represents the number of layers
+            need to be obfuscated of duplicate layers (such as transformer layers or resnet blocks).
+            Default: ``None``.
+        obf_config (dict): The configuration of model obfuscation polices. Default: ``None``.
         saved_path (str): The directory path for saving obfuscated ckpt files. Default: ``'./'``.
         obfuscate_scale (Union[float, int]): Obfuscate scale of weights. The generated random obf_ratios will be in
             range of (1 / obfuscate_scale, obfuscate_scale). Default: 100.
 
+    Returns:
+        dict[str], obf_metadata, which is the necessary data that needs to be load when running obfuscated network.
+
     Raises:
         TypeError: If `network` is not nn.Cell.
         TypeError: If `ckpt_files` is not string or `saved_path` is not string.
         TypeError: If `target_modules` is not list.
         TypeError: If target_modules's elements are not string.
+        TypeError: If obf_config is not dict.
         ValueError: If `ckpt_files` is not exist or `saved_path` is not exist.
         ValueError: If the number of elements of `target_modules` is less than ``2``.
         ValueError: If the first string of `target_modules` contains characters other than uppercase and lowercase
@@ -68,54 +160,91 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./', ob
         ValueError: If the third string of `target_modules` is not in the format of 'obfuscate_layers:all' or
             'obfuscate_layers:int'.
 
-    Returns:
-        list[float], obf_ratios, which is the necessary data that needs to be load when running obfuscated network.
-
     Examples:
         >>> from mindspore import obfuscate_ckpt, save_checkpoint
         >>> # Refer to https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> save_checkpoint(net, './test_net.ckpt')
         >>> target_modules = ['', 'fc1|fc2']
-        >>> obfuscate_ckpt(net, target_modules, './', './')
+        >>> obfuscate_ckpt(net, './', target_modules=target_modules, saved_path='./')
     """
+    def _gen_obfuscate_tensor(tensor_shape, tensor_type='rearrange'):
+        obf_tensor = None
+        if tensor_type == 'rearrange':
+            if len(tensor_shape) == 1:
+                obf_tensor = ms.Tensor(np.random.permutation(tensor_shape[0]), dtype=ms.int32)
+            if len(tensor_shape) == 2:
+                tensor = ms.Tensor(np.identity(tensor_shape[0]), dtype=ms.int32)
+                p = ms.Tensor(np.random.permutation(tensor_shape[1]), dtype=ms.int32)
+                obf_tensor = tensor[:, p]
+        if tensor_type == 'random':
+            obf_tensor = ms.Tensor(np.random.randint(1, obfuscate_scale, size=tensor_shape), dtype=ms.float16)
+        return obf_tensor
+
+    def _gen_obf_metadata(config):
+        name = config.get('name')
+        if name is None:
+            return False
+        save_metadata = config.get('save_metadata', False)
+        metadata_op_name = config.get('metadata_op')
+        layers = config.get('layers')
+        if not layers:
+            if not obf_metadata.get(name):
+                obf_tensor = _gen_obfuscate_tensor(config.get('shape'), config.get('type'))
+                obf_metadata[name] = obf_tensor
+                if save_metadata:
+                    saved_obf_tensor = obf_tensor
+                    if metadata_op_name is not None:
+                        metadata_op = _get_op(metadata_op_name)
+                        saved_obf_tensor = metadata_op(saved_obf_tensor)
+                    if saved_obf_tensor is not None:
+                        saved_metadata[name] = saved_obf_tensor.asnumpy()
+        else:
+            for layer in layers:
+                strTemplate = Template(name)
+                obf_name = strTemplate.safe_substitute({"layer": str(layer)})
+                obf_tensor = _gen_obfuscate_tensor(config.get('shape'), config.get('type'))
+                obf_metadata[obf_name] = obf_tensor
+                if save_metadata:
+                    saved_obf_tensor = obf_tensor
+                    if metadata_op_name is not None:
+                        metadata_op = _get_op(metadata_op_name)
+                        saved_obf_tensor = metadata_op(saved_obf_tensor)
+                    if saved_obf_tensor is not None:
+                        saved_metadata[obf_name] = saved_obf_tensor.asnumpy()
+        return True
+
     if not isinstance(network, nn.Cell):
         raise TypeError("network must be nn.Cell, but got {}.".format(type(network)))
     _check_dir_path('ckpt_files', ckpt_files)
     _check_dir_path('saved_path', saved_path)
-    # Try to find default target modules
-    if target_modules is None:
-        to_split_modules = _get_default_target_modules(ckpt_files)
-    else:
-        if len(target_modules) >= 1 and target_modules[0] == '/':
-            target_modules[0] = ''
-        to_split_modules = target_modules
-    if not _check_valid_target(network, to_split_modules):
-        raise ValueError("The obfuscate module path {} is not exist, please check the input 'target_modules'."
-                         .format(to_split_modules))
+
+    if obf_config is None:
+        if not _check_valid_target(network, target_modules):
+            raise ValueError("{} is not exist, please check the input 'target_modules'.".format(target_modules))
+        log.warning("'target_modules and obf_ratios' will be deprecated and "
+                    "removed in a future version, use 'obf_config' instead.")
+        obf_config = _transform_target_modules(target_modules)
+    if not isinstance(obf_config, dict):
+        raise TypeError("obf_config type should be dict, but got {}.".format(type(obf_config)))
+    if not obf_config or not _check_valid_obf_config(obf_config, 'obf_metadata_config')\
+        or not _check_valid_obf_config(obf_config, 'weight_obf_config'):
+        raise ValueError("'obf_config' is empty or not valid, please check the input.")
+    obf_metadata = {}
+    obf_metadata_config = obf_config.get('obf_metadata_config', [])
+    saved_metadata = {}
+    for config in obf_metadata_config:
+        _gen_obf_metadata(config)
     if (not isinstance(obfuscate_scale, (float, int))) or (obfuscate_scale <= 1):
         raise ValueError("obfuscate_scale must be float or int, and larger than 1, but got {}."
                          .format(obfuscate_scale))
-    # generate and save obf_ratios to saved_path
-    path_list = to_split_modules[0].split('/')
-    target_list = to_split_modules[1].split('|')
-    global OBF_RATIOS_LENGTH
-    number_of_ratios = OBF_RATIOS_LENGTH * OBF_RATIOS_WIDTH
-    if number_of_ratios > MAX_OBF_RATIOS_NUM:
-        OBF_RATIOS_LENGTH = MAX_OBF_RATIOS_NUM // OBF_RATIOS_WIDTH
-        number_of_ratios = OBF_RATIOS_LENGTH * OBF_RATIOS_WIDTH
-    obf_ratios = []
-    secrets_generator = secrets.SystemRandom()
-    for _ in range(number_of_ratios):
-        secure_float = secrets_generator.uniform(1 / obfuscate_scale, obfuscate_scale)
-        obf_ratios.append(secure_float)
     # start obfuscate ckpt
     ckpt_dir_files = os.listdir(ckpt_files)
     for ckpt_name in ckpt_dir_files:
-        sub_path = os.path.abspath(ckpt_files) + '/' + ckpt_name
+        sub_path = os.path.realpath(ckpt_files) + '/' + ckpt_name
         if Path(sub_path).is_dir():
             sub_ckpt_file_list = os.listdir(sub_path)
-            new_saved_path = os.path.abspath(saved_path) + '/' + ckpt_name
+            new_saved_path = os.path.realpath(saved_path) + '/' + ckpt_name
             if not os.path.exists(new_saved_path):
                 try:
                     os.mkdir(new_saved_path, mode=0o700)
@@ -124,71 +253,148 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./', ob
             for sub_ckpt_name in sub_ckpt_file_list:
                 if not sub_ckpt_name.endswith('.ckpt'):
                     continue
-                _obfuscate_single_ckpt(os.path.abspath(sub_path) + '/' + sub_ckpt_name, obf_ratios, path_list,
-                                       target_list, new_saved_path)
+                _obfuscate_single_ckpt(os.path.realpath(sub_path) + '/' + sub_ckpt_name, obf_metadata,
+                                       obf_config, new_saved_path)
         else:
             if not ckpt_name.endswith('.ckpt'):
                 continue
-            _obfuscate_single_ckpt(os.path.abspath(ckpt_files) + '/' + ckpt_name, obf_ratios, path_list,
-                                   target_list, saved_path)
-    return obf_ratios
+            _obfuscate_single_ckpt(os.path.realpath(ckpt_files) + '/' + ckpt_name,
+                                   obf_metadata, obf_config, saved_path)
+    return saved_metadata
 
 
-def _obfuscate_single_ckpt(ckpt_name, obf_ratios, path_list, target_list, saved_path):
+def _obfuscate_single_ckpt(ckpt_name, obf_metadata, obf_config, saved_path):
     """Obfuscate single ckpt file"""
-    module_has_been_obfuscated = set()
+    def _get_op_input_name(obf_op, name_key='input_x', layer=0):
+        op_name = obf_op.get('name')
+        input_name = obf_op.get(name_key)
+        if input_name is None:
+            log.error("can not find input: {} for op: {}.".format(name_key, op_name))
+            return None
+        strTemplate = Template(input_name)
+        input_name = strTemplate.safe_substitute({"layer": str(layer)})
+        return input_name
+
+    def _get_op_input(input_name, obf_param):
+        op_input = obf_metadata.get(input_name, None) if input_name.startswith('obf_metadata') else obf_param
+        return op_input
+
+    def _obfuscate_param(param, obf_metadata, obf_ops, layer=0):
+        param_dtype = F.dtype(param)
+        obf_param = param
+        for i in range(len(obf_ops)):
+            op_name = obf_ops[i].get('name')
+            if not isinstance(op_name, str):
+                raise TypeError('{} should be str type, but got {}'.format(op_name, type(op_name)))
+            if op_name == 'mul':
+                input_x = obf_param
+                input_y_name = _get_op_input_name(obf_ops[i], 'input_y', layer)
+                input_y = obf_metadata.get(input_y_name)
+                if input_x is None or input_y is None:
+                    log.error("input_x or input_y is None")
+                    return None
+                input_y = F.cast(input_y, param_dtype)
+                obf_param = ops.mul(input_x, input_y)
+            elif op_name == 'permuate':
+                input_x_name = _get_op_input_name(obf_ops[i], 'input_x', layer)
+                p = obf_metadata.get(input_x_name, None)
+                if p is None or obf_param is None:
+                    log.error("input_x or param is None")
+                    return None
+                obf_param = obf_param[p]
+            elif op_name == 'matmul':
+                input_x_name = _get_op_input_name(obf_ops[i], 'input_x', layer)
+                input_y_name = _get_op_input_name(obf_ops[i], 'input_y', layer)
+                input_x = _get_op_input(input_x_name, obf_param)
+                input_y = _get_op_input(input_y_name, obf_param)
+                if input_x is None or input_y is None:
+                    log.error("the input_x or input_y of op: {} is None.".format(op_name))
+                    return None
+                input_x = ops.transpose(input_x, (1, 0)) if obf_ops[i].get('transpose_a', False) else input_x
+                input_y = ops.transpose(input_y, (1, 0)) if obf_ops[i].get('transpose_b', False) else input_y
+                obf_param = ops.matmul(F.cast(input_x, param_dtype), F.cast(input_y, param_dtype))
+            else:
+                log.error("unsupported op, op must be matmul or permuate or mul, but got {}."
+                          .format(op_name))
+                return None
+        return obf_param
+
     try:
         ckpt_param = load_checkpoint(ckpt_name)
     except (ValueError, TypeError, OSError):
-        logger.error("Load checkpoint failed for file {}.".format(ckpt_name))
-        return None
-    obf_ratios_index = -1
+        log.error("Load checkpoint failed for file {}.".format(ckpt_name))
+        return False
+
+    weight_obf_config = obf_config.get('weight_obf_config', [])
     for item in ckpt_param:
-        module = _get_valid_module(item, path_list, target_list)
-        if module:
-            layer_index = _judge_layer_index(item)
-            if layer_index >= OBF_RATIOS_LENGTH:
-                continue
-            if module not in module_has_been_obfuscated:
-                module_has_been_obfuscated.add(module)
-                obf_ratios_index += 1
-            ratio_total_index = layer_index * OBF_RATIOS_WIDTH + obf_ratios_index % OBF_RATIOS_WIDTH
-            ckpt_param[item].set_data(ckpt_param[item].value() / obf_ratios[ratio_total_index])
+        item_split = item.split('.')
+        param_path = '/'.join(item_split[:len(item_split)])
+        for obf_target in weight_obf_config:
+            if not isinstance(obf_target, dict):
+                raise TypeError('{} should be dict type, but got {}'.format(obf_target, type(obf_target)))
+            target = obf_target.get('target', None)
+            layers = obf_target.get('layers', [])
+            obf_ops = obf_target.get('weight_obf_ops', None)
+            if not target or not obf_ops:
+                raise KeyError("target or obf_ops is None.")
+            if not layers:
+                if target == param_path:
+                    obf_param = _obfuscate_param(ckpt_param[item].value(), obf_metadata, obf_ops)
+                    if obf_param is None:
+                        log.error("obfuscate weight {} failed.".format(item))
+                        return False
+                    ckpt_param[item].set_data(obf_param)
+            for layer in layers:
+                strTemplate = Template(target)
+                target_path = strTemplate.safe_substitute({"layer": str(layer)})
+                if target_path == param_path:
+                    obf_param = _obfuscate_param(ckpt_param[item].value(), obf_metadata, obf_ops, layer)
+                    if obf_param is None:
+                        log.error("obfuscate weight {} failed.".format(item))
+                        return False
+                    ckpt_param[item].set_data(obf_param)
+
     # save the obfuscated model to saved_path
     obf_param_list = []
     for item in ckpt_param:
         obf_param_list.append({'name': item, 'data': ckpt_param[item]})
     ckpt_file_name = ckpt_name.split('/')[-1]
     obf_ckpt_file_name = ckpt_file_name.split('.')[0] + '_obf' + '.ckpt'
-    save_checkpoint(obf_param_list, os.path.abspath(saved_path) + '/' + obf_ckpt_file_name)
-    return None
+    save_checkpoint(obf_param_list, os.path.realpath(saved_path) + '/' + obf_ckpt_file_name)
+    return True
 
 
-def load_obf_params_into_net(network, target_modules, obf_ratios, data_parallel_num=1, **kwargs):
+def load_obf_params_into_net(network, target_modules=None, obf_ratios=None, obf_config=None,
+                             data_parallel_num=1, **kwargs):
     """
-    load obfuscate ratios into obfuscated network. Usually used in conjunction with :func:`mindspore.obfuscate_ckpt`
-    interface.
+    Modify model structure according to obfuscation config and load obfuscated checkpoint into obfuscated network.
 
     Args:
         network (nn.Cell): The original network that need to be obfuscated.
-        target_modules (list[str]): The target module of network that need to be obfuscated. The first string
-            represents the network path of target module in original network, which should be in form of ``'A/B/C'``.
-            The second string represents the obfuscation target module, which should be in form of ``'D|E|F'``. For
-            example, thr target_modules of GPT2 can be ``['backbone/blocks/attention', 'dense1|dense2|dense3']``.
-            If target_modules has the third value, it should be in the format of 'obfuscate_layers:all' or
-            'obfuscate_layers:int', which represents the number of layers need to be obfuscated of duplicate layers
-            (such as transformer layers or resnet blocks).
+        target_modules (list[str]): The target ops that need to be obfuscated in the network. The first string
+            represents the network path of the target ops in the original network, which should be in form of
+            ``"A/B/C"``. The second string represents the names of multiple target ops in the same path, which
+            should be in form of ``"D|E|F"``. For example, thr target_modules of GPT2 can be ``['backbone
+            /blocks/attention', 'dense1|dense2|dense3']``. If target_modules has the third value, it should be
+            in the format of 'obfuscate_layers:all' or 'obfuscate_layers:int', which represents the number of
+            layers need to be obfuscated of duplicate layers (such as transformer layers or resnet blocks).
+            Default: ``None``.
+        obf_ratios (Tensor): The obf ratios generated when execute :func:`mindspore.obfuscate_ckpt`. Default: ``None``.
+        obf_config (dict): The configuration of model obfuscation polices. Default: ``None``.
         data_parallel_num (int): The data parallel number of parallel training. Default: 1.
-        obf_ratios (Tensor): The obf ratios generated when execute :func:`mindspore.obfuscate_ckpt`.
         kwargs (dict): Configuration options dictionary.
 
             - ignored_func_decorators (list[str]): The name list of function decorators in network's python code.
             - ignored_class_decorators (list[str]): The name list of class decorators in network's python code.
 
+    Returns:
+        nn.Cell, new_net, which is the obfuscated network.
+
     Raises:
         TypeError: If `network` is not nn.Cell.
         TypeError: If `obf_ratios` is not Tensor.
         TypeError: If `target_modules` is not list.
+        TypeError: If `obf_config` is not dict.
         TypeError: If target_modules's elements are not string.
         ValueError: If the number of elements of `target_modules` is less than ``2``.
         ValueError: If `obf_ratios` is empty Tensor.
@@ -209,40 +415,33 @@ def load_obf_params_into_net(network, target_modules, obf_ratios, data_parallel_
         >>> save_checkpoint(net, './test_net.ckpt')
         >>> target_modules = ['', 'fc1|fc2']
         >>> # obfuscate ckpt files
-        >>> obfuscate_ckpt(net, target_modules, './', './')
+        >>> obfuscate_ckpt(net, './', target_modules=target_modules, saved_path='./')
         >>> # load obf ckpt into network
         >>> new_net = LeNet5()
         >>> load_checkpoint('./test_net_obf.ckpt', new_net)
-        >>> obf_ratios = Tensor(np.load('./obf_ratios.npy'), mstype.float16)
-        >>> obf_net = load_obf_params_into_net(new_net, target_modules, obf_ratios)
+        >>> obf_net = load_obf_params_into_net(new_net, target_modules)
     """
     if not isinstance(network, nn.Cell):
         raise TypeError("network must be nn.Cell, but got {}.".format(type(network)))
-    if not isinstance(obf_ratios, Tensor):
-        raise TypeError("obf_ratios must be MindSpore Tensor, but got {}.".format(type(obf_ratios)))
-    if obf_ratios.size == 0:
-        raise ValueError("obf_ratios can not be empty.")
-    if not _check_valid_target(network, target_modules):
-        raise ValueError("{} is not exist, please check the input 'target_modules'.".format(target_modules))
+    if obf_config is None:
+        if not _check_valid_target(network, target_modules):
+            raise ValueError("{} is not exist, please check the input 'target_modules'.".format(target_modules))
+        log.warning("'target_modules and obf_ratios' will be deprecated and "
+                    "removed in a future version, use 'obf_config' instead.")
+        obf_config = _transform_target_modules(target_modules)
+
+    if not isinstance(obf_config, dict):
+        raise TypeError('{} should be dict type, but got {}'.format(obf_config, type(obf_config)))
+
+    if not obf_config or not _check_valid_obf_config(obf_config, 'network_obf_config'):
+        raise ValueError("'obf_config' is empty or not valid, please check the input.")
+
     if (not isinstance(data_parallel_num, int)) or (data_parallel_num <= 0):
         raise ValueError("data_parallel_num must be positive number, but got {}.".format(data_parallel_num))
-    if len(target_modules) >= 1 and target_modules[0] == '/':
-        target_modules[0] = ''
-    path_list = target_modules[0].split('/')
-    path_len = len(path_list)
-    target_list = []
-    for _ in range(path_len):
-        target_list.append([])
-    target_list.append(target_modules[1].split('|'))
-    global MAX_OBF_RATIOS_NUM, OBF_RATIOS_LENGTH
-    number_of_ratios = OBF_RATIOS_LENGTH * OBF_RATIOS_WIDTH
-    if number_of_ratios > MAX_OBF_RATIOS_NUM:
-        OBF_RATIOS_LENGTH = MAX_OBF_RATIOS_NUM // OBF_RATIOS_WIDTH
-        number_of_ratios = OBF_RATIOS_LENGTH * OBF_RATIOS_WIDTH
-    MAX_OBF_RATIOS_NUM = number_of_ratios
-    rewrite_network = _obfuscate_network(network, path_list, target_list, data_parallel_num=data_parallel_num, **kwargs)
-    setattr(rewrite_network, 'obf_ratios', obf_ratios)
-    return rewrite_network
+
+    network_obf_config = obf_config.get('network_obf_config', [])
+    new_net = _obfuscate_network(network, network_obf_config, data_parallel_num=data_parallel_num, **kwargs)
+    return new_net
 
 
 def _check_dir_path(name, dir_path):
@@ -255,15 +454,6 @@ def _check_dir_path(name, dir_path):
         raise TypeError("{} must be a directory path, but got {}.".format(name, dir_path))
 
 
-def _judge_layer_index(layer_name):
-    """Judge the layer index of target layers"""
-    split_name = layer_name.split('.')
-    for split_str in split_name[:]:
-        if split_str.isdigit():
-            return int(split_str)
-    return 0
-
-
 def _check_valid_target(network, target_modules):
     """check whether the input 'target_modules' exists"""
     if not isinstance(target_modules, list):
@@ -314,7 +504,7 @@ def _check_valid_target(network, target_modules):
     OBF_RATIOS_WIDTH = 0
     for target in target_list:
         if not hasattr(net, target):
-            logger.warning("{} does not exist in the path {}".format(target, target_modules[0]))
+            log.warning("{} does not exist in the path {}".format(target, target_modules[0]))
         else:
             OBF_RATIOS_WIDTH += 1
     if OBF_RATIOS_WIDTH == 0:
@@ -323,6 +513,118 @@ def _check_valid_target(network, target_modules):
     return True
 
 
+def _check_ops_info(ops_info):
+    """check ops info config"""
+    for op in ops_info:
+        op_name = op.get('name')
+        if not isinstance(op_name, str):
+            raise TypeError("op_name type should be str, but got {}.".format(type(op_name)))
+        input_x_name = op.get('input_x')
+        if not isinstance(input_x_name, str):
+            raise TypeError("input_x_name type should be str, but got {}.".format(type(input_x_name)))
+        input_y_name = op.get('input_y')
+        if not isinstance(input_y_name, str):
+            raise TypeError("input_y_name type should be str, but got {}.".format(type(input_y_name)))
+        if not isinstance(op.get('transpose_a', False), bool):
+            raise TypeError("transpose_a type should be bool, but got {}.".format(type(op.get('transpose_a'))))
+        if not isinstance(op.get('transpose_b', False), bool):
+            raise TypeError("transpose_b type should be bool, but got {}.".format(type(op.get('transpose_b'))))
+
+
+def _check_new_input_info(insert_new_input):
+    """check new input config"""
+    if not isinstance(insert_new_input, list):
+        raise TypeError("obf_config[][]['insert_new_input'] type should be list, but got {}."
+                        .format(type(insert_new_input)))
+    for new_input in insert_new_input:
+        input_name = new_input.get('name')
+        if not isinstance(input_name, str):
+            raise TypeError("obf_config[][]['insert_new_input'][]['name'] type should be str, but got {}."
+                            .format(type(input_name)))
+
+
+def _check_obf_metadata_config(config):
+    """check obf metadata config"""
+    name = config.get('name')
+    if not name or not isinstance(name, str):
+        raise TypeError("obf_config[][]['name'] type should be str, but got {}.".format(type(name)))
+    shape = config.get('shape')
+    if not shape or not isinstance(shape, list):
+        raise TypeError("obf_config[][]['shape'] type should be list, but got {}.".format(type(shape)))
+    for item in shape:
+        if not isinstance(item, int):
+            raise TypeError("shape[] type should be int, but got {}.".format(type(item)))
+    save_metadata = config.get('save_metadata', True)
+    if not isinstance(save_metadata, bool):
+        raise TypeError("obf_config[][]['save_metadata'] type should be bool, but got {}."
+                        .format(type(save_metadata)))
+    metadata_type = config.get('type')
+    if metadata_type is not None:
+        if not isinstance(metadata_type, str) or metadata_type not in _supported_metadata_type:
+            raise TypeError("obf_config[][]['type'] should be str and must in {}, but got {}."
+                            .format(str(_supported_metadata_type), type(metadata_type)))
+
+
+def _check_weight_obf_config(config):
+    """check weight obfuscation config"""
+    target = config.get('target')
+    if not target or not isinstance(target, str):
+        raise TypeError("obf_config[][]['target'] type should be str, but got {}.".format(type(target)))
+    weight_obf_ops = config.get('weight_obf_ops', [])
+    if not isinstance(weight_obf_ops, list):
+        raise TypeError("obf_config[][]['weight_obf_ops'] type should be list, but got {}."
+                        .format(type(weight_obf_ops)))
+    _check_ops_info(weight_obf_ops)
+
+
+def _check_network_obf_config(config):
+    """check network obfuscation config"""
+    target = config.get('target')
+    if not target or not isinstance(target, str):
+        raise TypeError("obf_config[][]['target'] type should be str, but got {}.".format(type(target)))
+    module = config.get('module')
+    if not module or not isinstance(module, str):
+        raise TypeError("obf_config[][]['module'] type should be str, but got {}.".format(type(module)))
+    insert_new_input = config.get('insert_new_input', [])
+    _check_new_input_info(insert_new_input)
+    insert_ops = config.get('insert_ops', [])
+    if not isinstance(insert_ops, list):
+        raise TypeError("obf_config[][]['insert_ops'] type should be list, but got {}.".format(type(insert_ops)))
+    _check_ops_info(insert_ops)
+
+
+def _check_valid_obf_config(obf_config, config_type):
+    """check obfuscation config"""
+    if not isinstance(config_type, str) or config_type not in _supported_config_type:
+        raise TypeError("config_type must be str, and in {}, but got {}."
+                        .format(str(_supported_config_type), config_type))
+    for config_type_item in obf_config.keys():
+        if not isinstance(config_type_item, str) or config_type_item not in _supported_config_type:
+            raise TypeError("config_type must be str, and in {}, but got {}."
+                            .format(str(_supported_config_type), config_type_item))
+    config_list = obf_config.get(config_type)
+    if not isinstance(config_list, list):
+        raise TypeError("obf_config[] type of should be list, but got {}.".format(type(config_list)))
+
+    for config in config_list:
+        if not isinstance(config, dict):
+            raise TypeError("obf_config[][] type should be dict, but got {}.".format(type(config)))
+        if config_type == 'obf_metadata_config':
+            _check_obf_metadata_config(config)
+        elif config_type == 'weight_obf_config':
+            _check_weight_obf_config(config)
+        elif config_type == 'network_obf_config':
+            _check_network_obf_config(config)
+        layers = config.get('layers')
+        if layers is not None:
+            if not isinstance(layers, list):
+                raise TypeError("obf_config[][]['layers'] type should be list, but got {}.".format(type(layers)))
+            for layer in layers:
+                if not isinstance(layer, int):
+                    raise TypeError("obf_config[][]['layers'][] type should be int, but got {}.".format(type(layer)))
+    return True
+
+
 def _update_max_obf_ratios_num(target_modules):
     """Update MAX_OBF_RATIOS_NUM"""
     if len(target_modules) >= 3:
@@ -341,160 +643,163 @@ def _update_max_obf_ratios_num(target_modules):
             MAX_OBF_RATIOS_NUM = int(obfuscate_layers[1]) * OBF_RATIOS_WIDTH
 
 
-def _get_default_target_modules(ckpt_files):
-    """Get the default or suggested target modules, if the target modules is None."""
-
-    def _split_to_path_and_target(module, target):
-        # split module into path list and target list
-        target_index = module.index(target)
-        path = module[:target_index - 1]
-        target = module[target_index:].split('/')[0]
-        return path, target
-
-    def _find_default_obfuscate_modules(net_path):
-        # find modules including the default paths
-        default_module = {'attention'}
-        for module in default_module:
-            if module in net_path and module not in candidate_modules:
-                candidate_modules.append(net_path)
-        # find the default targets in the default module
-        default_target = {'dense', 'query', 'key', 'value'}
-        for target in default_target:
-            for candidate in candidate_modules:
-                if target in candidate:
-                    path, target = _split_to_path_and_target(candidate, target)
-                    if path not in paths:
-                        paths.append(path)
-                    if target not in targets:
-                        targets.append(target)
-
-    def _find_suggested_obfuscate_modules(net_path):
-        default_target = {'dense', 'query', 'key', 'value'}
-        for target in default_target:
-            # find the suggest modules
-            if target in net_path:
-                path, target = _split_to_path_and_target(net_path, target)
-                if [path, target] not in suggest_modules:
-                    suggest_modules.append([path, target])
-
-    # store the potential candidate_modules
-    candidate_modules = []
-    suggest_modules = []
-    paths = []
-    targets = []
-    ckpt_dir_files = os.listdir(ckpt_files)
-    for ckpt_name in ckpt_dir_files:
-        if not ckpt_name.endswith('.ckpt'):
-            continue
-        try:
-            ckpt_param = load_checkpoint(os.path.abspath(ckpt_files) + '/' + ckpt_name)
-        except (ValueError, TypeError, OSError):
-            logger.error("Load checkpoint failed for file {}.".format(os.path.abspath(ckpt_files) + '/' + ckpt_name))
-            return None
-        for item in ckpt_param:
-            param_path = _remove_digit(item)
-            param_path = '/'.join(param_path)
-            # find candidate modules including the default paths and append candidate_modules
-            _find_default_obfuscate_modules(param_path)
-            # give the suggested modules and find the default targets in the default module
-            _find_suggested_obfuscate_modules(param_path)
-    if paths and targets:
-        target_modules = [paths[0], '|'.join(targets)]
-        logger.warning("The default obfuscate modules is obtained:{}".format(target_modules))
-        return target_modules
-    # logging the suggested target module
-    logger.warning("The default obfuscate modules can not be obtained. The suggested possible paths are given below: {}"
-                   .format(suggest_modules))
-    raise ValueError("Can not get the default path, please specify the path in the form of ['A/B/C', 'D1|D2']")
-
-
-def _get_valid_module(item, path_list, target_list):
-    """get the valid module"""
-    number_path = len(path_list)
-    net_path = _remove_digit(item)
-    net_path = '/'.join(net_path[:number_path])
-    tar_path = '/'.join(path_list)
-    # update the weights with obf_ratios in target module
-    if net_path == tar_path:
-        for target in target_list:
-            if target in item.split('.'):
-                target_index = item.split('.').index(target)
-                module = ''.join(item.split('.')[:target_index + 1])
-                return module
-    return None
-
-
 def _remove_digit(item):
     """remove digit in the parameter path"""
-    param_path = item.split('.')
-    for tmp_str in param_path[:]:
+    item_split = item.split('_')
+    for tmp_str in item_split[:]:
         if tmp_str.isdigit():
-            param_path.remove(tmp_str)
-    return param_path
+            item_split.remove(tmp_str)
+    return '_'.join(item_split)
 
 
-def _obfuscate_network(model, path_list, target_list, data_parallel_num=1, **kwargs):
-    """obfuscate original network, including add mul operation and add inputs for passing obf_ratio."""
+def _remove_scope(item):
+    """remove scope of name values"""
+    item_split = item.split('.')
+    for tmp_str in item_split[:]:
+        if tmp_str == 'self':
+            item_split.remove(tmp_str)
+    return '.'.join(item_split)
 
-    def _insert_input(stree: SymbolTree, arg_name: str = 'y_obf'):
+
+def _obfuscate_network(model, obf_config=None, data_parallel_num=1, **kwargs):
+    """obfuscate original network, including add deobfuscation ops and add inputs for passing obf_metadata."""
+
+    def _insert_input(stree: SymbolTree, arg_name: str = 'obf_metadata'):
         """add inputs for passing obf_ratio"""
         last_input = None
         for node in stree.nodes():
             if node.get_node_type() == NodeType.Input:
                 last_input = node
         position = stree.after(last_input)
-        # the insert input node name would be 'input_y_obf'
+        # the insert input node name would be 'input_obf_metadata'
         new_input_node = last_input.create_input(arg_name)
         stree.insert(position, new_input_node)
 
-    def _insert_mul(stree: SymbolTree, node: Node, index: int):
-        """add mul operation for original network"""
-        arg_list = node.get_targets().copy()
-        input_y_node = stree.get_node("input_y_obf")
-        v: str = input_y_node.get_targets()[0].value
-        sv: ScopedValue = ScopedValue.create_naming_value(v + f'[{index}]')
-        arg_list.append(sv)
-        target_list = node.get_targets().copy()
-        if data_parallel_num > 1:
-            logger.info("Data parallel number is: {}".format(data_parallel_num))
-            new_mul_node = node.create_call_cell(cell=ops.Mul().shard(((data_parallel_num, 1), ())),
-                                                 targets=target_list, args=arg_list, name='mul')
+    def _update_subnet(substree: SymbolTree, subnode: Node):
+        """update the network once the subnet is obfuscated"""
+        input_y_node = substree.get_node("input_obf_metadata")
+        if input_y_node is None:
+            log.error("can not find input node: obf_metadata for net: {}.".format(subnode.get_name()))
+            return False
+        if hasattr(subnode, 'get_handler'):
+            subnode.get_handler().append_kwarg({"obf_metadata": input_y_node.get_targets()[0]})
         else:
-            new_mul_node = node.create_call_cell(cell=ops.Mul(), targets=target_list, args=arg_list, name='mul')
-        position = stree.after(node)
-        stree.insert(position, new_mul_node)
+            subnode.append_kwarg({"obf_metadata": input_y_node.get_targets()[0]})
+        return True
 
-    def _insert_mul_by_name(stree: SymbolTree, after_name_list: list):
+    def _insert_ops(stree: SymbolTree, node: Node, insert_ops: list):
+        """add mul operation for original network"""
+        current_node = node
+        for insert_op in insert_ops:
+            arg_list = current_node.get_targets().copy()
+            obf_metadata = stree.get_node("input_obf_metadata")
+            if obf_metadata is None:
+                raise ValueError("can not find input node: obf_metadata for net: {}.".format(current_node.get_name()))
+            v: str = obf_metadata.get_targets()[0].value
+            index = insert_op['input_y']
+            sv: ScopedValue = ScopedValue.create_naming_value(v + f'["{index}"]')
+            arg_list.append(sv)
+            target_list = current_node.get_targets().copy()
+            name = insert_op['name']
+            if data_parallel_num > 1:
+                new_node = current_node.create_call_cell(cell=_get_op(name).shard(((data_parallel_num, 1), ())),
+                                                         targets=target_list, args=arg_list, name=name)
+            else:
+                new_node = current_node.create_call_cell(cell=_get_op(name), targets=target_list, args=arg_list,
+                                                         name=name)
+            position = stree.after(current_node)
+            stree.insert(position, new_node)
+            current_node = new_node
+
+    def _insert_ops_by_name(stree: SymbolTree, after_name_list: list, module: str):
         """add mul operation after the target nodes according the name of them"""
         if not after_name_list:
             return
         for node in stree.nodes():
             for after_name in after_name_list:
                 if node.get_name() == after_name:
-                    global OBF_RATIOS_INSERT_INDEX
-                    if OBF_RATIOS_INSERT_INDEX < MAX_OBF_RATIOS_NUM:
-                        _insert_mul(stree, node, OBF_RATIOS_INSERT_INDEX)
-                        OBF_RATIOS_INSERT_INDEX += 1
-
-    def _update_subnet(substree: SymbolTree, subnode: Node):
-        """update the network once the subnet is obfuscated"""
-        input_y_node = substree.get_node("input_y_obf")
-        if input_y_node is None:
-            return
-        subnode.get_handler().append_kwarg({"y_obf": input_y_node.get_targets()[0]})
-
-    def _traverse(stree, i=0):
-        """traverse and obfuscate the original network"""
-        if len(path_list) == i:
-            return
+                    insert_ops = insert_ops_map[module+'/'+after_name]
+                    _insert_ops(stree, node, insert_ops)
+
+    def _process_controlflow_node(node: Node, stree: SymbolTree, full_path: str, path: str, targets: dict):
+        ctrl = node.get_handler() if hasattr(node, 'get_handler') else node
+        cell_loop_name = ''
+        find_cell_loop = False
+        if hasattr(ctrl, "loop_vars") and ctrl.loop_vars:
+            cell_loop_name = ctrl.loop_vars[0]
+            inputs = ctrl.get_inputs()
+            for input in inputs:
+                if input.get_node_type() == NodeType.CellContainer:
+                    find_cell_loop = True
+                    full_node_name = input.get_name()
+                    node_name = _remove_digit(_remove_scope(full_node_name))
+                    if not _process_cellcontainer_node(input, full_path+'/'+full_node_name,
+                                                       path+'/'+node_name, targets):
+                        log.error("_process_cellcontainer_node for node: {} failed.".format(node_name))
+                        return False
+        for c_node in ctrl.nodes():
+            c_node_name = c_node.get_name()
+            c_node_type = c_node.get_node_type()
+            if c_node.get_node_type() == NodeType.ControlFlow:
+                if not _process_controlflow_node(c_node, stree, full_path+'/'+c_node_name, path, targets):
+                    return False
+            elif c_node.get_node_type() == NodeType.Tree and _is_target_module(path + '/' + c_node_name, targets):
+                sub_stree = SymbolTree(c_node.symbol_tree)
+                _insert_input(sub_stree, arg_name='obf_metadata')
+                _insert_ops_by_name(sub_stree, after_name_list=targets.get(path + '/' + c_node_name, None),
+                                    module=path + '/' + c_node_name)
+                if not _traverse(sub_stree, full_path+'/'+c_node_name, path+'/'+c_node_name, targets):
+                    log.error("_traverse for node: {} failed.".format(c_node_name))
+                    return False
+                if not _update_subnet(sub_stree, c_node):
+                    log.error("_update_subnet for node: {} failed.".format(c_node_name))
+                    return False
+            elif find_cell_loop and c_node_type == NodeType.CallFunction and c_node_name.startswith(cell_loop_name):
+                input_y_node = stree.get_node("input_obf_metadata")
+                if input_y_node is None:
+                    log.error("input_y_node for node: {} is None.".format(c_node_name))
+                    return False
+                c_node.append_kwarg({"obf_metadata": input_y_node.get_targets()[0]})
+        return True
+
+    def _process_cellcontainer_node(node: Node, full_path: str, path: str, targets: dict):
+        cellcontainer = node.get_handler() if hasattr(node, 'get_handler') else node
+        for i in range(len(cellcontainer.nodes())):
+            cell_node = cellcontainer.nodes()[i]
+            # insert input for each sub_stree in cell_container
+            if _is_target_module(path, targets) and cell_node.get_node_type() == NodeType.Tree:
+                sub_stree = SymbolTree(cell_node.symbol_tree)
+                _insert_input(sub_stree, arg_name='obf_metadata')
+                _insert_ops_by_name(sub_stree, after_name_list=targets.get(path, None), module=path)
+                if not _traverse(sub_stree, full_path + '/' + str(i), path + '/' + str(i), targets):
+                    return False
+        return True
+
+    def _is_target_module(path, targets):
+        for target_module in targets.keys():
+            if target_module.startswith(path):
+                return True
+        return False
+
+    def _traverse(stree: SymbolTree, full_path: str, path: str, targets: dict):
         for node in stree.nodes():
             node_name = node.get_name()
-            if node.get_node_type() == NodeType.Tree and node_name.startswith(path_list[i]):
+            if node.get_node_type() == NodeType.ControlFlow:
+                if not _process_controlflow_node(node, stree, full_path + '/' + node_name, path, targets):
+                    log.error("process controlflow node: {} failed.".format(node.get_name()))
+                    return False
+            elif node.get_node_type() == NodeType.Tree and _is_target_module(path + '/' + node_name, targets):
                 sub_stree = node.get_sub_tree()
-                _traverse(sub_stree, i + 1)
-                _insert_input(sub_stree, arg_name='y_obf')
-                _insert_mul_by_name(sub_stree, after_name_list=target_list[i + 1])
-                _update_subnet(sub_stree, node)
+                _insert_input(sub_stree, arg_name='obf_metadata')
+                _insert_ops_by_name(sub_stree, after_name_list=targets.get(path + '/' + node_name, None),
+                                    module=path + '/' + node_name)
+                if not _traverse(sub_stree, full_path + '/' + node_name, path + '/' + node_name, targets):
+                    log.error("traverse sub_stree for node: {} failed.".format(node.get_name()))
+                    return False
+                if not _update_subnet(sub_stree, node):
+                    log.error("update subnet for node: {} failed.".format(node.get_name()))
+                    return False
+        return True
 
     def _register_denied_func_decorators(fn):
         """set the function decorators which should be denied for parse"""
@@ -523,9 +828,48 @@ def _obfuscate_network(model, path_list, target_list, data_parallel_num=1, **kwa
         if kw_class_dec and not isinstance(kw_class_dec[0], str):
             raise TypeError('elements of {} should be str, but got {}'.format(kw_class_dec, type(kw_class_dec[0])))
 
+    targets = {}
+    insert_ops_map = {}
+    for obf_item in obf_config:
+        module = obf_item.get('module', None)
+        target = obf_item.get('target', None)
+        insert_ops_info = obf_item.get('insert_ops', None)
+        layers = obf_item.get('layers', [])
+        if not layers:
+            real_insert_ops_info = []
+            if not targets.get(module, None):
+                targets[module] = []
+            if target not in targets[module]:
+                targets[module].append(target)
+            target_path = module + '/' + target
+            for op_info in insert_ops_info:
+                real_op_info = op_info.copy()
+                real_insert_ops_info.append(real_op_info)
+            insert_ops_map[target_path] = real_insert_ops_info
+        for layer in layers:
+            real_insert_ops_info = []
+            strTemplate = Template(module)
+            real_module = strTemplate.safe_substitute({"layer": str(layer)})
+            if not targets.get(real_module, None):
+                targets[real_module] = []
+            if target not in targets[real_module]:
+                targets[real_module].append(target)
+            target_path = real_module + '/' + target
+            for op_info in insert_ops_info:
+                real_op_info = op_info.copy()
+                strTemplate = Template(real_op_info['input_x'])
+                real_op_info['input_x'] = strTemplate.safe_substitute({"layer": str(layer)})
+                strTemplate = Template(real_op_info['input_y'])
+                real_op_info['input_y'] = strTemplate.safe_substitute({"layer": str(layer)})
+                real_insert_ops_info.append(real_op_info)
+            insert_ops_map[target_path] = real_insert_ops_info
+
+    root_path = ""
     main_stree = SymbolTree.create(model)
-    _traverse(main_stree, 0)
-    _insert_input(main_stree, arg_name='y_obf')
-    _insert_mul_by_name(main_stree, after_name_list=target_list[0])
+    _insert_input(main_stree, arg_name='obf_metadata')
+    _insert_ops_by_name(main_stree, after_name_list=targets.get(root_path, None), module=root_path)
+    if not _traverse(main_stree, full_path=root_path, path=root_path, targets=targets):
+        log.error("_traverse for root_path: {} failed.".format(root_path))
+        return None
     new_net = main_stree.get_network()
     return new_net