def extend_with_gradient_accumulation(base_optimizer, name=None): """返回新的优化器类,加入梯度累积 """ class new_optimizer(base_optimizer): """带有梯度累积的优化器 """ def __init__(self, grad_accum_steps, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.grad_accum_steps = grad_accum_steps self._first_get_gradients = True def get_gradients(self, loss, params): if self._first_get_gradients: self._first_get_gradients = False return super(new_optimizer, self).get_gradients(loss, params) else: return [ag / self.grad_accum_steps for ag in self.accum_grads] @K.symbolic def get_updates(self, loss, params): # 更新判据 cond = K.equal(self.iterations % self.grad_accum_steps, 0) cond = K.cast(cond, K.floatx()) # 获取梯度 grads = self.get_gradients(loss, params) self.accum_grads = [ K.zeros(K.int_shape(p), dtype=K.dtype(p), name='accum_grad_%s' % i) for i, p in enumerate(params) ] old_update = K.update def new_update(x, new_x): new_x = cond * new_x + (1 - cond) * x return old_update(x, new_x) K.update = new_update updates = super(new_optimizer, self).get_updates(loss, params) K.update = old_update # 累积梯度 with tf.control_dependencies(updates): accum_updates = [ K.update(ag, g + (1 - cond) * ag) for g, ag in zip(grads, self.accum_grads) ] return accum_updates def get_config(self): config = {'grad_accum_steps': self.grad_accum_steps} base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_piecewise_linear_lr_v2(base_optimizer, name=None): """返回新的优化器类,加入分段线性学习率 """ class new_optimizer(base_optimizer): """带有分段线性学习率的优化器 其中schedule是形如{1000: 1, 2000: 0.1}的字典, 表示0~1000步内学习率线性地从零增加到100%,然后 1000~2000步内线性地降到10%,2000步以后保持10% """ def __init__(self, lr_schedule, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.lr_schedule = {int(i): j for i, j in lr_schedule.items()} def _decayed_lr(self, var_dtype): lr_multiplier = piecewise_linear(self.iterations, self.lr_schedule) lr_t = super(new_optimizer, self)._decayed_lr(var_dtype) return lr_t * K.cast(lr_multiplier, var_dtype) def get_config(self): config = {'lr_schedule': self.lr_schedule} base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def encode(self, first_text, second_text=None, max_length=None, first_length=None, second_length=None): """输出文本对应token id和segment id 如果传入first_length,则强行padding第一个句子到指定长度; 同理,如果传入second_length,则强行padding第二个句子到指定长度。 """ if is_string(first_text): first_tokens = self.tokenize(first_text) else: first_tokens = first_text if second_text is None: second_tokens = None elif is_string(second_text): second_tokens = self.tokenize(second_text, add_cls=False) else: second_tokens = second_text if max_length is not None: self.truncate_sequence(max_length, first_tokens, second_tokens, -2) first_token_ids = self.tokens_to_ids(first_tokens) if first_length is not None: first_token_ids = first_token_ids[:first_length] first_token_ids.extend([self._token_pad_id] * (first_length - len(first_token_ids))) first_segment_ids = [0] * len(first_token_ids) if second_text is not None: second_token_ids = self.tokens_to_ids(second_tokens) if second_length is not None: second_token_ids = second_token_ids[:second_length] second_token_ids.extend( [self._token_pad_id] * (second_length - len(second_token_ids))) second_segment_ids = [1] * len(second_token_ids) first_token_ids.extend(second_token_ids) first_segment_ids.extend(second_segment_ids) return first_token_ids, first_segment_ids
def extend_with_lookahead(base_optimizer, name=None): """返回新的优化器类,加入look ahead """ class new_optimizer(base_optimizer): """带有look ahead的优化器 https://arxiv.org/abs/1907.08610 steps_per_slow_update: 即论文中的k; slow_step_size: 即论文中的alpha。 """ def __init__(self, steps_per_slow_update=5, slow_step_size=0.5, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.steps_per_slow_update = steps_per_slow_update self.slow_step_size = slow_step_size @K.symbolic def get_updates(self, loss, params): updates = super(new_optimizer, self).get_updates(loss, params) k, alpha = self.steps_per_slow_update, self.slow_step_size cond = K.equal(self.iterations % k, 0) slow_vars = [ K.zeros(K.int_shape(p), dtype=K.dtype(p), name='slow_var_%s' % i) for i, p in enumerate(params) ] with tf.control_dependencies(updates): slow_updates = [ K.update(q, K.switch(cond, q + alpha * (p - q), q)) for p, q in zip(params, slow_vars) ] with tf.control_dependencies(slow_updates): copy_updates = [ K.update(p, K.switch(cond, q, p)) for p, q in zip(params, slow_vars) ] return copy_updates def get_config(self): config = { 'steps_per_slow_update': self.steps_per_slow_update, 'slow_step_size': self.slow_step_size } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_lazy_optimization(base_optimizer, name=None): """返回新的优化器类,加入懒惰更新 """ class new_optimizer(base_optimizer): """带有懒惰更新的优化器 使得部分权重(尤其是embedding)只有在梯度不等于0时 才发生更新。 """ def __init__(self, include_in_lazy_optimization=None, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.include_in_lazy_optimization = include_in_lazy_optimization or [] self._first_get_gradients = True def get_gradients(self, loss, params): if self._first_get_gradients: self._first_get_gradients = False return super(new_optimizer, self).get_gradients(loss, params) else: return [self.grads[p] for p in params] @K.symbolic def get_updates(self, loss, params): self.grads = dict(zip(params, self.get_gradients(loss, params))) old_update = K.update def new_update(x, new_x): if is_one_of(x, params) and self._do_lazy_optimization(x): g = self.grads[x] r = K.any(K.not_equal(g, 0.), axis=-1, keepdims=True) new_x = x + (new_x - x) * K.cast(r, K.floatx()) return old_update(x, new_x) K.update = new_update updates = super(new_optimizer, self).get_updates(loss, params) K.update = old_update return updates def _do_lazy_optimization(self, w): return string_matching(w.name, self.include_in_lazy_optimization) def get_config(self): config = { 'include_in_lazy_optimization': self.include_in_lazy_optimization } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_layer_adaptation(base_optimizer, name=None): """返回新的优化器类,加入层自适应学习率 """ class new_optimizer(base_optimizer): """带有层自适应学习率的优化器 用每一层参数的模长来校正当前参数的学习率 https://arxiv.org/abs/1904.00962 """ def __init__(self, exclude_from_layer_adaptation=None, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.exclude_from_layer_adaptation = exclude_from_layer_adaptation or [] if not hasattr(self, 'learning_rate'): self.learning_rate = self.lr @K.symbolic def get_updates(self, loss, params): old_update = K.update def new_update(x, new_x): if is_one_of(x, params) and self._do_layer_adaptation(x): dx = new_x - x lr_t = K.clip(self.learning_rate, K.epsilon(), 1e10) x_norm = tf.norm(x) g_norm = tf.norm(dx / lr_t) ratio = K.switch( x_norm > 0., K.switch(g_norm > K.epsilon(), x_norm / g_norm, 1.), 1.) new_x = x + dx * ratio return old_update(x, new_x) K.update = new_update updates = super(new_optimizer, self).get_updates(loss, params) K.update = old_update return updates def _do_layer_adaptation(self, w): return (not string_matching(w.name, self.exclude_from_layer_adaptation)) def get_config(self): config = { 'exclude_from_layer_adaptation': self.exclude_from_layer_adaptation } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_layer_adaptation_v2(base_optimizer, name=None): """返回新的优化器类,加入层自适应学习率 """ class new_optimizer(base_optimizer): """带有层自适应学习率的优化器 用每一层参数的模长来校正当前参数的学习率 https://arxiv.org/abs/1904.00962 """ def __init__(self, exclude_from_layer_adaptation=None, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.exclude_from_layer_adaptation = exclude_from_layer_adaptation or [] def _resource_apply_op(self, grad, var, indices=None): old_update = K.update def new_update(x, new_x): if x is var and self._do_layer_adaptation(x): dx = new_x - x lr_t = self._decayed_lr(x.dtype.base_dtype) lr_t = K.clip(lr_t, K.epsilon(), 1e10) x_norm = tf.norm(x) g_norm = tf.norm(dx / lr_t) ratio = K.switch( x_norm > 0., K.switch(g_norm > K.epsilon(), x_norm / g_norm, 1.), 1.) new_x = x + dx * ratio return old_update(x, new_x) K.update = new_update op = super(new_optimizer, self)._resource_apply_op(grad, var, indices) K.update = old_update return op def _do_layer_adaptation(self, w): return (not string_matching(w.name, self.exclude_from_layer_adaptation)) def get_config(self): config = { 'exclude_from_layer_adaptation': self.exclude_from_layer_adaptation } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_lookahead_v2(base_optimizer, name=None): """返回新的优化器类,加入look ahead """ class new_optimizer(base_optimizer): """带有look ahead的优化器 https://arxiv.org/abs/1907.08610 steps_per_slow_update: 即论文中的k; slow_step_size: 即论文中的alpha。 """ def __init__(self, steps_per_slow_update=5, slow_step_size=0.5, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.steps_per_slow_update = steps_per_slow_update self.slow_step_size = slow_step_size def _create_slots(self, var_list): super(new_optimizer, self)._create_slots(var_list) for var in var_list: self.add_slot(var, 'slow_var') def _resource_apply_op(self, grad, var, indices=None): op = super(new_optimizer, self)._resource_apply_op(grad, var, indices) k, alpha = self.steps_per_slow_update, self.slow_step_size cond = K.equal(self.iterations % k, 0) slow_var = self.get_slot(var, 'slow_var') slow_var_t = slow_var + alpha * (var - slow_var) with tf.control_dependencies([op]): slow_update = K.update(slow_var, K.switch(cond, slow_var_t, slow_var)) with tf.control_dependencies([slow_update]): copy_update = K.update(var, K.switch(cond, slow_var, var)) return copy_update def get_config(self): config = { 'steps_per_slow_update': self.steps_per_slow_update, 'slow_step_size': self.slow_step_size } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_gradient_accumulation_v2(base_optimizer, name=None): """返回新的优化器类,加入梯度累积 """ class new_optimizer(base_optimizer): """带有梯度累积的优化器 """ def __init__(self, grad_accum_steps, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.grad_accum_steps = grad_accum_steps def _create_slots(self, var_list): super(new_optimizer, self)._create_slots(var_list) for var in var_list: self.add_slot(var, 'ag') 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 get_config(self): config = {'grad_accum_steps': self.grad_accum_steps} base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_lazy_optimization_v2(base_optimizer, name=None): """返回新的优化器类,加入懒惰更新 """ class new_optimizer(base_optimizer): """带有懒惰更新的优化器 使得部分权重(尤其是embedding)只有在梯度不等于0时 才发生更新。 """ def __init__(self, include_in_lazy_optimization=None, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.include_in_lazy_optimization = include_in_lazy_optimization or [] def _resource_apply_op(self, grad, var, indices=None): old_update = K.update def new_update(x, new_x): if x is var and self._do_lazy_optimization(x): if indices is None: r = K.any(K.not_equal(grad, 0.), axis=-1, keepdims=True) new_x = x + (new_x - x) * K.cast(r, K.floatx()) return old_update(x, new_x) else: return self._resource_scatter_add( x, indices, K.gather(new_x - x, indices)) return old_update(x, new_x) K.update = new_update op = super(new_optimizer, self)._resource_apply_op(grad, var, indices) K.update = old_update return op def _do_lazy_optimization(self, w): return string_matching(w.name, self.include_in_lazy_optimization) def get_config(self): config = { 'include_in_lazy_optimization': self.include_in_lazy_optimization } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_weight_decay(base_optimizer, name=None): """返回新的优化器类,加入权重衰减 """ class new_optimizer(base_optimizer): """带有权重衰减的优化器 """ def __init__(self, weight_decay_rate, exclude_from_weight_decay=None, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.weight_decay_rate = weight_decay_rate self.exclude_from_weight_decay = exclude_from_weight_decay or [] if not hasattr(self, 'learning_rate'): self.learning_rate = self.lr @K.symbolic def get_updates(self, loss, params): old_update = K.update def new_update(x, new_x): if is_one_of(x, params) and self._do_weight_decay(x): new_x = new_x - self.learning_rate * self.weight_decay_rate * x return old_update(x, new_x) K.update = new_update updates = super(new_optimizer, self).get_updates(loss, params) K.update = old_update return updates def _do_weight_decay(self, w): return (not string_matching(w.name, self.exclude_from_weight_decay)) def get_config(self): config = { 'weight_decay_rate': self.weight_decay_rate, 'exclude_from_weight_decay': self.exclude_from_weight_decay } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def extend_with_weight_decay_v2(base_optimizer, name=None): """返回新的优化器类,加入权重衰减 """ class new_optimizer(base_optimizer): """带有权重衰减的优化器 """ def __init__(self, weight_decay_rate, exclude_from_weight_decay=None, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.weight_decay_rate = weight_decay_rate self.exclude_from_weight_decay = exclude_from_weight_decay or [] def _resource_apply_op(self, grad, var, indices=None): old_update = K.update def new_update(x, new_x): if x is var and self._do_weight_decay(x): lr_t = self._decayed_lr(x.dtype.base_dtype) new_x = new_x - lr_t * self.weight_decay_rate * x return old_update(x, new_x) K.update = new_update op = super(new_optimizer, self)._resource_apply_op(grad, var, indices) K.update = old_update return op def _do_weight_decay(self, w): return (not string_matching(w.name, self.exclude_from_weight_decay)) def get_config(self): config = { 'weight_decay_rate': self.weight_decay_rate, 'exclude_from_weight_decay': self.exclude_from_weight_decay } base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def __init__(self, token_dict, do_lower_case=False): """初始化 """ super(Tokenizer, self).__init__(do_lower_case) if is_string(token_dict): token_dict = load_vocab(token_dict) self._token_dict = token_dict self._token_dict_inv = {v: k for k, v in token_dict.items()} for token in ['pad', 'cls', 'sep', 'unk', 'mask']: try: _token_id = token_dict[getattr(self, '_token_%s' % token)] setattr(self, '_token_%s_id' % token, _token_id) except: pass self._vocab_size = len(token_dict)
def extend_with_piecewise_linear_lr(base_optimizer, name=None): """返回新的优化器类,加入分段线性学习率 """ class new_optimizer(base_optimizer): """带有分段线性学习率的优化器 其中schedule是形如{1000: 1, 2000: 0.1}的字典, 表示0~1000步内学习率线性地从零增加到100%,然后 1000~2000步内线性地降到10%,2000步以后保持10% """ def __init__(self, lr_schedule, *args, **kwargs): super(new_optimizer, self).__init__(*args, **kwargs) self.lr_schedule = {int(i): j for i, j in lr_schedule.items()} @K.symbolic def get_updates(self, loss, params): lr_multiplier = piecewise_linear(self.iterations, self.lr_schedule) old_update = K.update def new_update(x, new_x): if is_one_of(x, params): new_x = x + (new_x - x) * lr_multiplier return old_update(x, new_x) K.update = new_update updates = super(new_optimizer, self).get_updates(loss, params) K.update = old_update return updates def get_config(self): config = {'lr_schedule': self.lr_schedule} base_config = super(new_optimizer, self).get_config() return dict(list(base_config.items()) + list(config.items())) if is_string(name): new_optimizer.__name__ = name keras.utils.get_custom_objects()[name] = new_optimizer return new_optimizer
def call(self, inputs, mask=None, a_mask=None, p_bias=None): """实现多头注意力 q_mask: 对输入的query序列的mask。 主要是将输出结果的padding部分置0。 v_mask: 对输入的value序列的mask。 主要是防止attention读取到padding信息。 a_mask: 对attention矩阵的mask。 不同的attention mask对应不同的应用。 p_bias: 在attention里的位置偏置。 一般用来指定相对位置编码的种类。 """ q, k, v = inputs[:3] q_mask, v_mask, n = None, None, 3 if mask is not None: if mask[0] is not None: q_mask = K.cast(mask[0], K.floatx()) if mask[2] is not None: v_mask = K.cast(mask[2], K.floatx()) if a_mask: if len(inputs) == 3: a_mask = 'history_only' else: a_mask = inputs[n] n += 1 # 线性变换 qw = self.q_dense(q) kw = self.k_dense(k) vw = self.v_dense(v) # 形状变换 qw = K.reshape(qw, (-1, K.shape(q)[1], self.heads, self.key_size)) kw = K.reshape(kw, (-1, K.shape(k)[1], self.heads, self.key_size)) vw = K.reshape(vw, (-1, K.shape(v)[1], self.heads, self.head_size)) # Attention a = tf.einsum('bjhd,bkhd->bhjk', qw, kw) # 处理位置编码 if p_bias == 'typical_relative': pos_embeddings = inputs[n] a = a + tf.einsum('bjhd,jkd->bhjk', qw, pos_embeddings) elif p_bias == 't5_relative': pos_embeddings = K.permute_dimensions(inputs[n], (2, 0, 1)) a = a + K.expand_dims(pos_embeddings, 0) # Attention(续) if p_bias != 't5_relative': # T5不用缩放 a = a / self.key_size**0.5 a = sequence_masking(a, v_mask, 1, -1) if a_mask is not None: if is_string(a_mask): ones = K.ones_like(a[:1, :1]) a_mask = (ones - tf.linalg.band_part(ones, -1, 0)) * 1e12 a = a - a_mask else: a = a - (1 - a_mask) * 1e12 a = K.softmax(a) # 完成输出 o = tf.einsum('bhjk,bkhd->bjhd', a, vw) if p_bias == 'typical_relative': o = o + tf.einsum('bhjk,jkd->bjhd', a, pos_embeddings) o = K.reshape(o, (-1, K.shape(o)[1], self.out_dim)) o = self.o_dense(o) # 返回结果 o = sequence_masking(o, q_mask, 0) return o