def _resource_apply_op(self, grad, var, indices=None): # 更新判据 cond = K.equal(self.iterations % self.grad_accum_steps, 0) # 获取梯度 ag = self.get_slot(var, 'ag') old_update = K.update def new_update(x, new_x): new_x = K.switch(cond, new_x, x) return old_update(x, new_x) K.update = new_update ag_t = ag / self.grad_accum_steps op = super(new_optimizer, self)._resource_apply_op(ag_t, var) K.update = old_update # 累积梯度 with tf.control_dependencies([op]): ag_t = K.switch(cond, K.zeros_like(ag), ag) with tf.control_dependencies([K.update(ag, ag_t)]): if indices is None: ag_t = K.update(ag, ag + grad) else: ag_t = self._resource_scatter_add(ag, indices, grad) return ag_t
def adversarial_training(model, embedding_name, epsilon=1): """给模型添加对抗训练 其中model是需要添加对抗训练的keras模型,embedding_name 则是model里边Embedding层的名字。要在模型compile之后使用。 """ if model.train_function is None: # 如果还没有训练函数 model._make_train_function() # 手动make old_train_function = model.train_function # 备份旧的训练函数 # 查找Embedding层 for output in model.outputs: embedding_layer = search_layer(output, embedding_name) if embedding_layer is not None: break if embedding_layer is None: raise Exception('Embedding layer not found') # 求Embedding梯度 embeddings = embedding_layer.embeddings # Embedding矩阵 gradients = K.gradients(model.total_loss, [embeddings]) # Embedding梯度 gradients = K.zeros_like(embeddings) + gradients[0] # 转为dense tensor # 封装为函数 inputs = (model._feed_inputs + model._feed_targets + model._feed_sample_weights) # 所有输入层 embedding_gradients = K.function( inputs=inputs, outputs=[gradients], name='embedding_gradients', ) # 封装为函数 def train_function(inputs): # 重新定义训练函数 grads = embedding_gradients(inputs)[0] # Embedding梯度 delta = epsilon * grads / (np.sqrt((grads**2).sum()) + 1e-8) # 计算扰动 K.set_value(embeddings, K.eval(embeddings) + delta) # 注入扰动 outputs = old_train_function(inputs) # 梯度下降 K.set_value(embeddings, K.eval(embeddings) - delta) # 删除扰动 return outputs model.train_function = train_function # 覆盖原训练函数