{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Context\n\nThis notebook demonstrates the use of **an algorithm to encrypt any AI model's parameters. The aim is to degrade the model's performance for unauthorised users** without a secret key. The algorithm was proposed in the paper [AdvParams: An Active DNN Intellectual Property Protection Technique via Adversarial Perturbation Based Parameter Encryption](https://arxiv.org/abs/2105.13697) by Xue et al.\n\nFor more details and documentation on reusing this algorithm, please see [this Github repository](https://github.com/Madhav-Malhotra/ML-parameter-encryption).\n\n------------","metadata":{}},{"cell_type":"markdown","source":"# Setup\n\nLoad libraries, the encryption dataset, and the model.","metadata":{}},{"cell_type":"code","source":"import os\nimport random\nimport numpy as np\nimport tensorflow as tf\nfrom datetime import datetime\nfrom tensorflow.keras import preprocessing\n\ndev = None\nif tf.test.is_gpu_available():\n    dev = \"/gpu:0\"\nelse:\n    dev = \"/cpu:0\"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-05-05T14:02:43.748249Z","iopub.execute_input":"2023-05-05T14:02:43.749459Z","iopub.status.idle":"2023-05-05T14:02:53.553384Z","shell.execute_reply.started":"2023-05-05T14:02:43.749401Z","shell.execute_reply":"2023-05-05T14:02:53.551876Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Helper functions to convert images to processed tensors\ndef convert_img(img):\n    img = preprocessing.image.resize(img, [256, 256])\n    img = preprocessing.image.crop_center(img, 224, 224)\n    img = tf.convert_to_tensor(img, dtype=tf.float32)\n    img = tf.transpose(img, perm=[1, 2, 0])\n    return img\n\ndef normalise_img(img_tensor):\n    mean = [0.485, 0.456, 0.406]\n    std = [0.229, 0.224, 0.225]\n    img_tensor = (img_tensor - mean) / std\n    return img_tensor\n\n\ndef load_imgs(root_dir : str, max_per_class : int):\n    '''\n    Loads images into a matrix of preprocessed tensors.\n    \n    Parameters\n    -------------------\n    root_dir (type: string)\n    - A filepath to a folder with subfolders with images. \n    - Each subfolder corresponds to one class.\n    max_per_class (type: int)\n    - The maximum number of images to save from each class.\n    \n    Returns\n    -------------------\n    tf.Tensor\n    - Shape is num_images x num_channels x width x height\n    '''\n    \n    images = []\n    labels = []\n    y = 0\n    \n    # Go through sub-class folders\n    for subclass in os.listdir(root_dir):\n        folder_name = os.path.join(root_dir, subclass)\n        files = os.listdir(folder_name)\n        print('.', end=\"\")\n        \n        # Go through up to max_per_class files in each sub-class\n        i = 0\n        max_files = min(len(files), max_per_class)\n        while (i < max_files):\n            \n            # Save valid files\n            file_name = os.path.join(folder_name, files[i])\n            if os.path.isfile(file_name):\n                converted = convert_img(Image.open(file_name))\n                if (converted.shape.as_list() == [3, 224, 224]):\n                    images.append(normalise_img(converted))\n                    labels.append(y)\n                    \n            i += 1\n        y += 1\n        \n    # return output\n    with tf.device(dev):\n        encrypt_imgs = tf.stack(images, axis=0)\n        encrypt_labels = tf.constant(labels, dtype=tf.int16)\n    return encrypt_imgs, encrypt_labels ","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:02:53.556882Z","iopub.execute_input":"2023-05-05T14:02:53.558665Z","iopub.status.idle":"2023-05-05T14:02:53.57404Z","shell.execute_reply.started":"2023-05-05T14:02:53.558609Z","shell.execute_reply":"2023-05-05T14:02:53.572483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_data(load_mode : bool=True, imgs_per_class : int=1) -> list:\n    '''\n    Returns image set for encryption (images, labels)\n    \n    Parameters\n    ------------------\n    load_mode (type: bool)\n    - Whether to load saved data or process new data\n    imgs_per_class (type: int)\n    - Number of images to load for each class\n    \n    Returns\n    encrypt_imgs (type: tf.Tensor, dim: N x C x H x W)\n    - Preprocessed images to use for encryption\n    encrypt_labels (type: tf.Tensor, dim: N)\n    - Labels for above images\n    '''\n    \n    # Load saved tensors\n    with tf.device(dev):\n        if (load_mode):\n            encrypt_imgs = tf.constant(\n                np.load('/kaggle/input/tempsampleilsvrc/encrypt_imgs.npy'),\n                dtype=tf.float32\n            )\n            # Need to reshape channels to be last dim for pretrained model\n            encrypt_imgs = tf.transpose(encrypt_imgs, perm=[0, 2, 3, 1])\n\n            encrypt_labels = tf.constant(\n                np.load('/kaggle/input/tempsampleilsvrc/encrypt_labels.npy'),\n                dtype=tf.int32\n            )\n\n        # Process new tensors and then save\n        else:\n            encrypt_imgs, encrypt_labels = load_imgs(\n                '/kaggle/input/imagenet-object-localization-challenge/ILSVRC/Data/CLS-LOC/train', \n                imgs_per_class)\n\n            f = gzip.GzipFile(\"encrypt_imgs.npy.gz\", \"w\")\n            numpy.save(file=f, arr=encrypt_imgs.detach().cpu().numpy())\n            f.close()\n\n            f = gzip.GzipFile(\"encrypt_labels.npy.gz\", \"w\")\n            numpy.save(file=f, arr=encrypt_labels.detach().cpu().numpy())\n            f.close()\n\n    return encrypt_imgs, encrypt_labels","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:02:53.576321Z","iopub.execute_input":"2023-05-05T14:02:53.578105Z","iopub.status.idle":"2023-05-05T14:02:53.593506Z","shell.execute_reply.started":"2023-05-05T14:02:53.578053Z","shell.execute_reply":"2023-05-05T14:02:53.591761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load data\nencrypt_imgs, encrypt_labels = get_data()\n\n# Load MobileNetV2\nmodel = tf.keras.applications.mobilenet_v2.MobileNetV2(\n    input_shape=None,\n    alpha=1.0,\n    weights='imagenet',\n    classifier_activation=None) # we want to get back model logits\nmodel.trainable = True","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:02:53.597526Z","iopub.execute_input":"2023-05-05T14:02:53.598709Z","iopub.status.idle":"2023-05-05T14:03:04.357889Z","shell.execute_reply.started":"2023-05-05T14:02:53.598659Z","shell.execute_reply":"2023-05-05T14:03:04.356568Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Encryption Functions","metadata":{}},{"cell_type":"code","source":"def get_layer_set(max_layers: int, model: tf.Module) -> list:\n    '''\n    Randomly selects model layers to choose parameters for encryption\n    Parameters\n    --------------------\n    max_layers (type: int)\n    - The maximum number of layers to be selected in the encryption set\n    model (type: tf.Module - or a derivative class of this)\n    - The model from which to draw layers from\n    Returns\n    --------------------\n    list\n    - Has names of the selected layers\n    '''\n\n    # Arrange layers into list\n    random.seed(datetime.now().timestamp())\n    model.trainable = True\n    selected_layers = [layer.name for layer in model.trainable_variables]\n\n    # Prune layers randomly if too many selected\n    while len(selected_layers) > max_layers:\n        rand_i = random.randint(0, len(selected_layers) - 1)\n        del selected_layers[rand_i]\n\n    return selected_layers","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:03:04.359452Z","iopub.execute_input":"2023-05-05T14:03:04.36065Z","iopub.status.idle":"2023-05-05T14:03:04.370004Z","shell.execute_reply.started":"2023-05-05T14:03:04.360611Z","shell.execute_reply":"2023-05-05T14:03:04.368338Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class EncryptionUtils:\n    def __init__(\n        self, \n        model : tf.Module, \n        encrypt_data : tf.data.Dataset, \n        encrypt_layers : list, \n        max_per_layer : int, \n        loss_threshold : float, \n        step_size : float, \n        boundary_distance : float,\n        device : str):\n        \n        self.model = model\n        self.encrypt_data = encrypt_data\n        self.encrypt_layers = encrypt_layers\n        self.max_per_layer = max_per_layer\n        self.loss_threshold = loss_threshold\n        self.step_size = step_size\n        self.boundary_distance = boundary_distance\n        self.device = device\n    \n    def compute_bounds(self, layer_params : tf.Tensor) -> list:\n        '''\n        Computes boundaries for parameter values in selected layer\n        \n        Parameters\n        --------------------\n        layer_params (type: tf.Tensor, dim: variable)\n        - Weights (multidimensional) or biases (1D) of current layer.\n        \n        Returns\n        --------------------\n        bound_low (type: float)\n        - lower limit on parameter values\n        bound_high (type: float)\n        - upper limit on parameter values\n        '''\n        \n        l_max = float(tf.math.reduce_max(layer_params))\n        l_min = float(tf.math.reduce_min(layer_params))\n        l_range = self.boundary_distance * (l_max - l_min)\n\n        return l_min+l_range, l_max-l_range\n    \n    def compute_loss_grad(self, param_name):\n        '''\n        Computes average loss and gradient across minibatches\n        \n        Parameters\n        -----------------------\n        param_name (type: str)\n        - Has an identifier for the trianable variable to be updated. \n        \n        Returns\n        -----------------------\n        avg_loss (type: tf.Tensor, dim: 0, dtype: float)\n        - Average loss across all batches\n        avg_grad (type: tf.Tensor, dim: variable, dtype: float)\n        - Average gradient across all batches. \n        '''\n        \n        # prep variables\n        avg_loss = 0\n        batches = 0\n        avg_grad = None\n        \n        trainable_var = None\n        for var in self.model.trainable_variables:\n            if var.name == param_name:\n                trainable_var = [var]\n                break\n        \n        \n        # go through batches\n        for x, y in self.encrypt_data:\n            print('.', end='')\n            batches += 1\n            \n            # Compute loss\n            with tf.GradientTape() as tape:\n                pred = self.model(x)\n                # Unlike pytorch, returns one loss value per tensor in batch\n                loss = tf.keras.losses.sparse_categorical_crossentropy(y, pred, from_logits=True)\n            avg_loss += tf.math.reduce_mean(loss) # Averages loss across batch\n            \n            # Compute gradient\n            dloss_dparams = tape.gradient(loss, trainable_var)            \n            if avg_grad == None: \n                avg_grad = dloss_dparams[0]\n            else:\n                avg_grad += dloss_dparams[0]\n            \n        return avg_loss / batches, avg_grad / batches\n\n    def get_max_grad(self, mask, grad):\n        '''\n        Finds the maximum gradient vector component\n        \n        Parameters\n        --------------------\n        mask (type: tf.Tensor, dim: variable)\n        - A tensor of 1s/0s for each parameter in current layer \n        - 1 indicates a parameter is updatable. \n        - Has same dim as current layer\n        grad (type: tf.Tensor, dim: variable)\n        - Gradient of current layer\n        \n        Returns\n        --------------------\n        val (type: float)\n        - Maximum gradient component for current layer's params\n        i (type: int)\n        - Max gradient component's index for current layer's params\n        - Layers with multi-dimensional parameters are flattened\n          before index is computed.\n        '''\n        \n        grad = grad * mask\n        val = float(tf.reduce_max(grad))\n        # returns flattened index for multidimensional tensors\n        i = tf.argmax(tf.reshape(grad, [-1]))\n\n        return val, int(i)\n    \n    def compute_update(self, param, grad, param_index):\n        '''\n        Computes update step for selected parameter\n        \n        Parameters\n        --------------------\n        param (type: tf.Tensor, dim: variable)\n        - Layer housing selected parameter\n        grad (type: float)\n        - Gradient component for selected parameter\n        param_index (type: int)\n        - Index of selected parameter within the layer\n        \n        Returns\n        --------------------\n        update (type: float)\n        -  new value for selected parameter\n        unrolled_i (type: tuple)\n        - the index of the element to update\n        '''\n        \n        # compute update step\n        param_range = tf.math.reduce_max(param) - tf.math.reduce_min(param)\n        step = self.step_size * int(grad > 0) * param_range\n\n        # get new param val\n        unrolled_i = self.unroll_index(param_index, param.shape)\n        # https://www.tensorflow.org/guide/tensor_slicing\n        new_val = tf.slice(param, begin=unrolled_i, size=[1 for i in unrolled_i]) + step\n        return new_val, unrolled_i\n    \n    \n    def unroll_index(self, i, new_dim):\n        '''\n        Turns a one-dimensional index into a new dimensional shape\n        \n        Parameters\n        --------------------\n        i (type: int)\n        - index of flattened tensor\n        new_dim (type: tuple)\n        - dimensions to adjust the flattened index into\n        \n        Returns\n        --------------------\n        tuple\n        - 1+ elements representing indices along different dimensions\n        '''\n        \n        # Handle exception for low-dimensional index\n        if len(new_dim) == 1:\n            return [i]\n        elif not len(new_dim):\n            return [0]\n        \n        # Create useful variables\n        out_i = []\n        new_dim = list(new_dim)\n        prod = int(new_dim[0])\n        for j in range(1, len(new_dim)):\n            prod *= int(new_dim[j])\n        \n        while (len(new_dim)):\n            # Get current index\n            prod /= int(new_dim[0])\n            out_i.append(int(i // prod))\n            \n            # Prepare for next round\n            i = int(i % prod)\n            del new_dim[0]\n        \n        return tuple(out_i)\n                        \n    def encrypt_parameters(self):\n        '''\n        Makes selectively-targeted adversarial modifications to parameters.\n\n        Parameters\n        --------------------\n        None\n\n        Returns: \n        ---------------------\n        dict\n        - has keys for each encrypted layer and values indicating their \n        adjusted weights\n        '''\n        \n        # Init variables and settings\n        modifications = {}\n        for layer in self.model.layers: \n            layer.trainable = False\n\n        # Go through each model layer\n        for layer in self.model.layers:\n            param = None\n            \n            # Check if layer has variables to encrypt\n            for v in layer.variables:\n                if v.name in self.encrypt_layers:\n                    layer.trainable = True\n                    param = v\n                    break\n            \n            # Skip to next layer if no layers to encrypt\n            if (type(param) == type(None)): \n                continue\n            print(param.name)\n            \n            # Generate mask to track which params in each layer are unusable\n            modifications[param.name] = []\n            mask = tf.ones(param.numpy().shape)\n            bound_low, bound_high = self.compute_bounds(param.read_value())\n            \n            # Only run updates on current layer up to max iters\n            for i in range(self.max_per_layer):\n                print(i, end=\"\")\n                avg_loss, avg_grad = self.compute_loss_grad(param.name)\n\n                # Abort if loss raised sufficiently\n                if avg_loss > self.loss_threshold:\n                    return modifications\n\n                # Get adversarial update\n                val, i = self.get_max_grad(mask, avg_grad)\n                new_val, unrolled_i = self.compute_update(param.read_value(), val, i)\n                \n                # Mask current parameter if its gradients are too large\n                if float(new_val) < bound_low or float(new_val) > bound_high:\n                    # https://www.tensorflow.org/api_docs/python/tf/tensor_scatter_nd_update\n                    mask = tf.tensor_scatter_nd_update(mask, [unrolled_i], tf.constant([0.0]))\n                \n                # Save diff between old value and new value\n                clipped = float(tf.clip_by_value(new_val, clip_value_min=bound_low, clip_value_max=bound_high))\n                current = float(tf.slice(param.read_value(), begin=unrolled_i, size=[1 for i in unrolled_i]))\n                modifications[param.name].append((unrolled_i, clipped - current))\n                \n                # Update old value\n                updated_params = param.read_value()\n                updated_params = tf.tensor_scatter_nd_update(updated_params, [unrolled_i], tf.constant([clipped]))\n                param.assign(updated_params)\n            \n            # Disable grad after layer operations finished\n            layer.trainable = False\n            \n        return modifications","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:04:36.325286Z","iopub.execute_input":"2023-05-05T14:04:36.326572Z","iopub.status.idle":"2023-05-05T14:04:36.357692Z","shell.execute_reply.started":"2023-05-05T14:04:36.32653Z","shell.execute_reply":"2023-05-05T14:04:36.356483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Encrypt Model Parameters","metadata":{}},{"cell_type":"code","source":"# Init hyperparameters\ndataset = tf.data.Dataset.from_tensor_slices((encrypt_imgs, encrypt_labels))\nencrypt_data = dataset.batch(64)\nencrypt_layers = get_layer_set(25, model) # max 25 layers adjusted\n\nmax_per_layer = 25         # max weights adjusted per layer\nstep_size = 0.1            # for gradient updates\n\n# boundary distance hyperparameter must be 0-0.5\n# lower bound: 0 = encrypted weight values can be anywhere in range of current weight values\n# upper bound: 0.5 = encrypted weight values go at the midpoint of current weight values\nboundary_distance = 0.1    \n\n# Test current loss before encryption\nx,y = next(iter(encrypt_data))\nloss = tf.keras.losses.sparse_categorical_crossentropy(y, model(x), from_logits=True)\nloss = tf.math.reduce_mean(loss)\nprint(\"Loss before: \", loss)\nloss_threshold = float(loss) * 5","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:04:40.359985Z","iopub.execute_input":"2023-05-05T14:04:40.360962Z","iopub.status.idle":"2023-05-05T14:04:42.84026Z","shell.execute_reply.started":"2023-05-05T14:04:40.360919Z","shell.execute_reply":"2023-05-05T14:04:42.838993Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"instance = EncryptionUtils(model, encrypt_data, encrypt_layers, max_per_layer, loss_threshold, step_size, boundary_distance, dev)","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:03:33.560248Z","iopub.execute_input":"2023-05-05T14:03:33.560686Z","iopub.status.idle":"2023-05-05T14:03:33.566791Z","shell.execute_reply.started":"2023-05-05T14:03:33.560648Z","shell.execute_reply":"2023-05-05T14:03:33.565417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"modifications = instance.encrypt_parameters()","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:03:36.819717Z","iopub.execute_input":"2023-05-05T14:03:36.820125Z","iopub.status.idle":"2023-05-05T14:03:40.995316Z","shell.execute_reply.started":"2023-05-05T14:03:36.820092Z","shell.execute_reply":"2023-05-05T14:03:40.994274Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Show loss after encryption\nx,y = next(iter(encrypt_data))\nloss = tf.keras.losses.sparse_categorical_crossentropy(y, model(x), from_logits=True)\nloss = tf.math.reduce_mean(loss)\nprint(\"Encrypted loss: \", loss)","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:03:45.86295Z","iopub.execute_input":"2023-05-05T14:03:45.863399Z","iopub.status.idle":"2023-05-05T14:03:46.681345Z","shell.execute_reply.started":"2023-05-05T14:03:45.86336Z","shell.execute_reply":"2023-05-05T14:03:46.680126Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Decrypt Parameters","metadata":{}},{"cell_type":"code","source":"def decrypt_parameters(model : tf.Module, modifications : dict) -> None:\n    '''\n    Decrypts parameters using secret key\n    \n    Parameters\n    --------------------\n    model (type: tf.Module - or a derivative class)\n    - The model containing parameters to decrypt\n    modifications (type: dict)\n    - An object recording the modifications made to parameters by layer.\n    \n    Returns\n    --------------------\n    None\n    '''\n    \n    encrypted_layers = modifications.keys()\n    \n    # Select model layer\n    for param in model.variables:\n        if param.name in encrypted_layers:\n            print(param.name)\n            \n            # Replace modified parameter\n            for index, value in modifications[param.name]:\n                updated_params = param.read_value()\n                current = float(tf.slice(updated_params, begin=index, size=[1 for i in index]))\n                updated_params = tf.tensor_scatter_nd_update(updated_params, [index], tf.constant([current - value]))\n                param.assign(updated_params)","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:03:55.394629Z","iopub.execute_input":"2023-05-05T14:03:55.395091Z","iopub.status.idle":"2023-05-05T14:03:55.403873Z","shell.execute_reply.started":"2023-05-05T14:03:55.39505Z","shell.execute_reply":"2023-05-05T14:03:55.402592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"decrypt_parameters(model, modifications)","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:03:59.395609Z","iopub.execute_input":"2023-05-05T14:03:59.396033Z","iopub.status.idle":"2023-05-05T14:03:59.406404Z","shell.execute_reply.started":"2023-05-05T14:03:59.395995Z","shell.execute_reply":"2023-05-05T14:03:59.405161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test loss after decryption\nx,y = next(iter(encrypt_data))\nloss = tf.keras.losses.sparse_categorical_crossentropy(y, model(x), from_logits=True)\nloss = tf.math.reduce_mean(loss)\nprint(\"Decrypted: \", loss)","metadata":{"execution":{"iopub.status.busy":"2023-05-05T14:04:01.912562Z","iopub.execute_input":"2023-05-05T14:04:01.913525Z","iopub.status.idle":"2023-05-05T14:04:02.736662Z","shell.execute_reply.started":"2023-05-05T14:04:01.913481Z","shell.execute_reply":"2023-05-05T14:04:02.735483Z"},"trusted":true},"execution_count":null,"outputs":[]}]}