def average_n_grads(split_grads):
    split_grads = [grad for grad in split_grads if grad is not None]

    # Dealing with IndexedSlices for large-dimensional embedding
    # matrix. The gradient of an embedding matrix is not a tensor,
    # but a tuple-like object named `IndexedSlices`, for this one,
    # we need to take special processings.
    if split_grads[0].__str__().startswith('IndexedSlices'):
        all_values = [grad.values for grad in split_grads]

        values = tf.divide(tf.add_n(all_values), len(split_grads))
        indices = split_grads[0].indices
        dense_shape = split_grads[0].dense_shape

        return tf.IndexedSlices(
            values=values,
            indices=indices,
            dense_shape=dense_shape)
    return tf.divide(tf.add_n(split_grads), len(split_grads))

r = tf.IndexedSlices(values=r,
                     indices=grad.indices,
                     dense_shape=grad.dense_shape)

init_r = tf.get_variable(
    'init_r',
    shape=[module.batch_size * module.max_seq_length,
           param.shape.as_list()[-1]],
    initializer=tf.random_uniform_initializer(
        minval=-epsilon, maxval=epsilon),
    trainable=False)

r = tf.IndexedSlices(values=r,
                     indices=grad.indices,
                     dense_shape=grad.dense_shape)

InvalidArgumentError (see above for traceback): data.shape = [4096,768] does not start with segment_ids.shape = [1024]
	 [[node add_1/y (defined at /root/unif-tencent/uf/processing.py:590)  = UnsortedSegmentSum[T=DT_FLOAT, Tindices=DT_INT32, Tnumsegments=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](truediv_202, bert/embeddings/Reshape/_457, add_1/strided_slice)]]
	 [[{{node Assign/_476}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_2147_Assign", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

    def _freelb(self, module, alpha=0.3, epsilon=0.3, n_loop=3, **kwargs):
        # FreeLB is similar to PGD, but uses average gradients from loop.
        # i.e. grad = (first_grad + ... + last_grad) / n_loop
        #
        # Also, it initializes the perturbation not from usual forward
        # propagation, but a collection of uniform distribution within
        # epsilon range. It does not uses actual gradient to average
        # gradients. The perturbation is iterated, in the same way with
        #  PGD.
        # (epsilon: the norm of perturbation, must be smaller than the
        # norm of gradients)

        # initialize
        (d_grads, module._losses, module._probs, module._preds) = \
            module._parallel_forward(**self._kwargs)
        grad, param = utils.get_grad_and_param(
            module.trainable_variables, d_grads, 'word_embedding')
        init_r = tf.get_variable(
            'init_r',
            shape=[module.batch_size * module.max_seq_length,
                   param.shape.as_list()[-1]],
            initializer=tf.random_uniform_initializer(
                minval=-epsilon, maxval=epsilon),
            trainable=False)
        init_op = tf.variables_initializer([init_r])
        with tf.control_dependencies([init_op]):    # fix perturbation
            # Scale randomly initialized permutation, to make sure norm
            # of `r` is smaller than epsilon.
            shape = tf.cast(np.prod(init_r.shape.as_list()), tf.float32)
            r = tf.divide(init_r, tf.sqrt(shape))
            r = tf.IndexedSlices(values=r,
                                 indices=grad.indices,
                                 dense_shape=grad.dense_shape)

        # with tf.control_dependencies([init_op]):    # fix perturbation
        #     # Scale randomly initialized permutation, to make sure norm
        #     # of `r` is smaller than epsilon.
        #     r = tf.divide(init_r, tf.norm(init_r, np.inf))
        #     r = tf.IndexedSlices(values=r,
        #                          indices=grad.indices,
        #                          dense_shape=grad.dense_shape)
        #     attack_op = param.assign(param + r)

        # attack
        acc_r = r
        all_grads = []
        for k in range(n_loop):
            attack_param = param + acc_r  ######修改部分
            attack_param = tf.stop_gradient(attack_param)  ######修改部分
            module.attack_trainable_variables = [attack_param if v.name == 'bert/embeddings/word_embeddings:0' else v for v in
                                                 module.trainable_variables]  ######修改部分
            (attack_grads, _, _, _) = \
                module._parallel_forward(attack_embeddings=attack_param, **self._kwargs)  ######修改部分
            all_grads.append(attack_grads)
            grad, _ = utils.get_grad_and_param(
                module.attack_trainable_variables,
                attack_grads, attack_param.name)
            tmp_r = tf.multiply(alpha, grad / (tf.norm(grad) + 1e-9))

            # In order not to shuffle the distribution of gradient-
            # induced perturbation, we use norm to scale instead of
            # simply clip the values.
            norm = tf.norm(acc_r + tmp_r)
            cur_r = tf.cond(norm > epsilon,
                            lambda: (acc_r + tmp_r) * tf.divide(epsilon, norm),
                            lambda: (acc_r + tmp_r))
            acc_r = cur_r

        attack_param = param + acc_r  ######修改部分
        attack_param = tf.stop_gradient(attack_param)  ######修改部分
        module.attack_trainable_variables = [attack_param if v.name == 'bert/embeddings/word_embeddings:0' else v for v in
                                             module.trainable_variables]  ######修改部分
        (attack_grads, _, _, _) = \
            module._parallel_forward(attack_embeddings=attack_param, **self._kwargs)  ######修改部分
        all_grads.append(attack_grads)

        # sum up
        grads = [utils.average_n_grads(split_grad) for split_grad in zip(
            *all_grads)]
        update_params_op = utils.update_global_params(
            module.trainable_variables, module._global_step,
            module._optimizer, grads)
        update_step_op = module._global_step.assign(module._global_step + 1)
        module._train_op = tf.group([update_params_op, update_step_op])

InvalidArgumentError (see above for traceback): Input to reshape is a tensor with 1024 values, but the requested shape has 4096
         [[node gradients_4/bert_4/embeddings/embedding_look_up_grad/Reshape_1 (defined at /jizhi/jizhi2/worker/trainer/uf/core.py:859)  = Reshape[T=DT_INT32, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](bert_16/embeddings/ExpandDims, gradients_16/bert_16/embeddings/embedding_look_up_grad/ExpandDims)]]
         [[{{node concat_2/_10269}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_124285_concat_2", tensor_type=DT_INT64, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

class FreeAT(tf.keras.Model):
    def train_step(self, data):
        x, y = data
        last_r = 0.0
        last_r_slice = 0.0
        K = 3
        ep = 1e-3

        for t in range(K):
            with tf.GradientTape() as tape:
                y_pred = self(x, training=True)
                loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)        
            embedding_gradients = tape.gradient(loss, [self.trainable_variables[0]])[0]
            grad_values = tf.zeros_like(self.trainable_variables[0]) + embedding_gradients
            sign = tf.cast(tf.greater(grad_values, 0.0), tf.float32)
            r = last_r + tf.multiply(ep, sign) if t > 0 else \
                    tf.multiply(ep, sign)
            r *= tf.divide(ep, tf.norm(r))
            r_slice = tf.IndexedSlices(
                values=r,
                indices=embedding_gradients.indices,
                dense_shape=embedding_gradients.dense_shape)
            self.trainable_variables[0].assign_add(r_slice - last_r_slice)
            last_r = r
            last_r_slice = r_slice

        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
        self.compiled_metrics.update_state(y, y_pred)
        
        return {m.name: m.result() for m in self.metrics}