forked from mindspore/mindspore
306 lines
11 KiB
Python
306 lines
11 KiB
Python
import mindspore as ms
|
|
import mindspore.nn as nn
|
|
import mindspore.ops as ops
|
|
import numpy as np
|
|
from math import sqrt
|
|
from utils.masking import TriangularCausalMask, ProbMask
|
|
|
|
class FullAttention(nn.Cell):
|
|
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
|
super(FullAttention, self).__init__()
|
|
self.scale = scale
|
|
self.mask_flag = mask_flag
|
|
self.output_attention = output_attention
|
|
self.dropout = nn.Dropout(attention_dropout)
|
|
|
|
def construct(self, queries, keys, values, attn_mask):
|
|
B, L, H, E = queries.shape
|
|
_, S, _, D = values.shape
|
|
scale = self.scale or 1. / sqrt(E)
|
|
|
|
scores = ops.einsum("blhe,bshe->bhls", queries, keys)
|
|
if self.mask_flag:
|
|
if attn_mask is None:
|
|
attn_mask = TriangularCausalMask(B, L)
|
|
|
|
scores = ops.masked_fill(scores, attn_mask.mask, -np.inf)
|
|
|
|
A = self.dropout(ops.softmax(scale * scores, axis=-1))
|
|
V = ops.einsum("bhls,bshd->blhd", A, values)
|
|
|
|
if self.output_attention:
|
|
return (V, A)
|
|
else:
|
|
return (V, None)
|
|
|
|
class ProbAttention(nn.Cell):
|
|
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
|
super(ProbAttention, self).__init__()
|
|
self.factor = factor
|
|
self.scale = scale
|
|
self.mask_flag = mask_flag
|
|
self.output_attention = output_attention
|
|
self.dropout = nn.Dropout(attention_dropout)
|
|
|
|
def _prob_QK(self, Q, K, sample_k, n_top):
|
|
B, H, L_K, E = K.shape
|
|
_, _, L_Q, _ = Q.shape
|
|
|
|
K_expand = K.expand_dims(-3).repeat(1, 1, L_Q, 1)
|
|
index_sample = ms.Tensor(np.random.randint(0, L_K, (L_Q, sample_k)), ms.int32)
|
|
K_sample = K_expand[:, :, ms.Tensor(np.arange(L_Q)).expand_dims(1), index_sample, :]
|
|
Q_K_sample = ops.matmul(Q.expand_dims(-2), K_sample.transpose(0, 1, 3, 2)).squeeze(-2)
|
|
|
|
M = Q_K_sample.max(-1)[0] - ops.reduce_sum(Q_K_sample, -1) / L_K
|
|
M_top = M.topk(n_top)[1]
|
|
|
|
Q_reduce = Q[ms.Tensor(np.arange(B))[:, None, None],
|
|
ms.Tensor(np.arange(H))[None, :, None],
|
|
M_top, :]
|
|
Q_K = ops.matmul(Q_reduce, K.transpose(0, 1, 3, 2))
|
|
|
|
return Q_K, M_top
|
|
|
|
def _get_initial_context(self, V, L_Q):
|
|
B, H, L_V, D = V.shape
|
|
if not self.mask_flag:
|
|
V_sum = ops.reduce_mean(V, axis=-2)
|
|
context = V_sum.expand_dims(-2).repeat(1, 1, L_Q, 1)
|
|
else:
|
|
assert(L_Q == L_V)
|
|
context = ops.cumsum(V, axis=-2)
|
|
return context
|
|
|
|
def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
|
|
B, H, L_V, D = V.shape
|
|
|
|
if self.mask_flag:
|
|
attn_mask = ProbMask(B, H, L_Q, index, scores)
|
|
scores = ops.masked_fill(scores, attn_mask.mask, -np.inf)
|
|
|
|
attn = ops.softmax(scores, axis=-1)
|
|
|
|
context_in[ms.Tensor(np.arange(B))[:, None, None],
|
|
ms.Tensor(np.arange(H))[None, :, None],
|
|
index, :] = ops.matmul(attn, V).astype(context_in.dtype)
|
|
|
|
if self.output_attention:
|
|
attns = (ms.Tensor(np.ones([B, H, L_V, L_V]) / L_V, dtype=context_in.dtype))
|
|
attns[ms.Tensor(np.arange(B))[:, None, None], ms.Tensor(np.arange(H))[None, :, None], index, :] = attn
|
|
return (context_in, attns)
|
|
else:
|
|
return (context_in, None)
|
|
|
|
def construct(self, queries, keys, values, attn_mask):
|
|
B, L_Q, H, D = queries.shape
|
|
_, L_K, _, _ = keys.shape
|
|
|
|
queries = ops.transpose(queries, (0, 2, 1, 3))
|
|
keys = ops.transpose(keys, (0, 2, 1, 3))
|
|
values = ops.transpose(values, (0, 2, 1, 3))
|
|
|
|
U_part = self.factor * int(np.ceil(np.log(L_K))) # c*ln(L_k)
|
|
u = self.factor * int(np.ceil(np.log(L_Q))) # c*ln(L_q)
|
|
|
|
U_part = min(U_part, L_K)
|
|
u = min(u, L_Q)
|
|
|
|
scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)
|
|
|
|
scale = self.scale or 1. / sqrt(D)
|
|
if scale is not None:
|
|
scores_top = scores_top * scale
|
|
|
|
context = self._get_initial_context(values, L_Q)
|
|
context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)
|
|
|
|
return ops.transpose(context, (0, 2, 1, 3)), attn
|
|
|
|
class AttentionLayer(nn.Cell):
|
|
def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None, mix=False):
|
|
super(AttentionLayer, self).__init__()
|
|
|
|
d_keys = d_keys or (d_model // n_heads)
|
|
d_values = d_values or (d_model // n_heads)
|
|
|
|
self.inner_attention = attention
|
|
self.query_projection = nn.Dense(d_model, d_keys * n_heads)
|
|
self.key_projection = nn.Dense(d_model, d_keys * n_heads)
|
|
self.value_projection = nn.Dense(d_model, d_values * n_heads)
|
|
self.out_projection = nn.Dense(d_values * n_heads, d_model)
|
|
self.n_heads = n_heads
|
|
self.mix = mix
|
|
|
|
def construct(self, queries, keys, values, attn_mask):
|
|
B, L, _ = queries.shape
|
|
_, S, _ = keys.shape
|
|
H = self.n_heads
|
|
|
|
queries = self.query_projection(queries).view(B, L, H, -1)
|
|
keys = self.key_projection(keys).view(B, S, H, -1)
|
|
values = self.value_projection(values).view(B, S, H, -1)
|
|
|
|
out, attn = self.inner_attention(
|
|
queries,
|
|
keys,
|
|
values,
|
|
attn_mask
|
|
)
|
|
if self.mix:
|
|
out = ops.transpose(out, (0, 2, 1, 3)).contiguous()
|
|
out = out.view(B, L, -1)
|
|
|
|
return self.out_projection(out), attn
|
|
import mindspore as ms
|
|
import mindspore.nn as nn
|
|
import mindspore.ops as ops
|
|
import numpy as np
|
|
from math import sqrt
|
|
from utils.masking import TriangularCausalMask, ProbMask
|
|
|
|
class FullAttention(nn.Cell):
|
|
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
|
super(FullAttention, self).__init__()
|
|
self.scale = scale
|
|
self.mask_flag = mask_flag
|
|
self.output_attention = output_attention
|
|
self.dropout = nn.Dropout(attention_dropout)
|
|
|
|
def construct(self, queries, keys, values, attn_mask):
|
|
B, L, H, E = queries.shape
|
|
_, S, _, D = values.shape
|
|
scale = self.scale or 1. / sqrt(E)
|
|
|
|
scores = ops.einsum("blhe,bshe->bhls", queries, keys)
|
|
if self.mask_flag:
|
|
if attn_mask is None:
|
|
attn_mask = TriangularCausalMask(B, L)
|
|
|
|
scores = ops.masked_fill(scores, attn_mask.mask, -np.inf)
|
|
|
|
A = self.dropout(ops.softmax(scale * scores, axis=-1))
|
|
V = ops.einsum("bhls,bshd->blhd", A, values)
|
|
|
|
if self.output_attention:
|
|
return (V, A)
|
|
else:
|
|
return (V, None)
|
|
|
|
class ProbAttention(nn.Cell):
|
|
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
|
super(ProbAttention, self).__init__()
|
|
self.factor = factor
|
|
self.scale = scale
|
|
self.mask_flag = mask_flag
|
|
self.output_attention = output_attention
|
|
self.dropout = nn.Dropout(attention_dropout)
|
|
|
|
def _prob_QK(self, Q, K, sample_k, n_top):
|
|
B, H, L_K, E = K.shape
|
|
_, _, L_Q, _ = Q.shape
|
|
|
|
K_expand = K.expand_dims(-3).repeat(1, 1, L_Q, 1)
|
|
index_sample = ms.Tensor(np.random.randint(0, L_K, (L_Q, sample_k)), ms.int32)
|
|
K_sample = K_expand[:, :, ms.Tensor(np.arange(L_Q)).expand_dims(1), index_sample, :]
|
|
Q_K_sample = ops.matmul(Q.expand_dims(-2), K_sample.transpose(0, 1, 3, 2)).squeeze(-2)
|
|
|
|
M = Q_K_sample.max(-1)[0] - ops.reduce_sum(Q_K_sample, -1) / L_K
|
|
M_top = M.topk(n_top)[1]
|
|
|
|
Q_reduce = Q[ms.Tensor(np.arange(B))[:, None, None],
|
|
ms.Tensor(np.arange(H))[None, :, None],
|
|
M_top, :]
|
|
Q_K = ops.matmul(Q_reduce, K.transpose(0, 1, 3, 2))
|
|
|
|
return Q_K, M_top
|
|
|
|
def _get_initial_context(self, V, L_Q):
|
|
B, H, L_V, D = V.shape
|
|
if not self.mask_flag:
|
|
V_sum = ops.reduce_mean(V, axis=-2)
|
|
context = V_sum.expand_dims(-2).repeat(1, 1, L_Q, 1)
|
|
else:
|
|
assert(L_Q == L_V)
|
|
context = ops.cumsum(V, axis=-2)
|
|
return context
|
|
|
|
def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
|
|
B, H, L_V, D = V.shape
|
|
|
|
if self.mask_flag:
|
|
attn_mask = ProbMask(B, H, L_Q, index, scores)
|
|
scores = ops.masked_fill(scores, attn_mask.mask, -np.inf)
|
|
|
|
attn = ops.softmax(scores, axis=-1)
|
|
|
|
context_in[ms.Tensor(np.arange(B))[:, None, None],
|
|
ms.Tensor(np.arange(H))[None, :, None],
|
|
index, :] = ops.matmul(attn, V).astype(context_in.dtype)
|
|
|
|
if self.output_attention:
|
|
attns = (ms.Tensor(np.ones([B, H, L_V, L_V]) / L_V, dtype=context_in.dtype))
|
|
attns[ms.Tensor(np.arange(B))[:, None, None], ms.Tensor(np.arange(H))[None, :, None], index, :] = attn
|
|
return (context_in, attns)
|
|
else:
|
|
return (context_in, None)
|
|
|
|
def construct(self, queries, keys, values, attn_mask):
|
|
B, L_Q, H, D = queries.shape
|
|
_, L_K, _, _ = keys.shape
|
|
|
|
queries = ops.transpose(queries, (0, 2, 1, 3))
|
|
keys = ops.transpose(keys, (0, 2, 1, 3))
|
|
values = ops.transpose(values, (0, 2, 1, 3))
|
|
|
|
U_part = self.factor * int(np.ceil(np.log(L_K))) # c*ln(L_k)
|
|
u = self.factor * int(np.ceil(np.log(L_Q))) # c*ln(L_q)
|
|
|
|
U_part = min(U_part, L_K)
|
|
u = min(u, L_Q)
|
|
|
|
scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)
|
|
|
|
scale = self.scale or 1. / sqrt(D)
|
|
if scale is not None:
|
|
scores_top = scores_top * scale
|
|
|
|
context = self._get_initial_context(values, L_Q)
|
|
context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)
|
|
|
|
return ops.transpose(context, (0, 2, 1, 3)), attn
|
|
|
|
class AttentionLayer(nn.Cell):
|
|
def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None, mix=False):
|
|
super(AttentionLayer, self).__init__()
|
|
|
|
d_keys = d_keys or (d_model // n_heads)
|
|
d_values = d_values or (d_model // n_heads)
|
|
|
|
self.inner_attention = attention
|
|
self.query_projection = nn.Dense(d_model, d_keys * n_heads)
|
|
self.key_projection = nn.Dense(d_model, d_keys * n_heads)
|
|
self.value_projection = nn.Dense(d_model, d_values * n_heads)
|
|
self.out_projection = nn.Dense(d_values * n_heads, d_model)
|
|
self.n_heads = n_heads
|
|
self.mix = mix
|
|
|
|
def construct(self, queries, keys, values, attn_mask):
|
|
B, L, _ = queries.shape
|
|
_, S, _ = keys.shape
|
|
H = self.n_heads
|
|
|
|
queries = self.query_projection(queries).view(B, L, H, -1)
|
|
keys = self.key_projection(keys).view(B, S, H, -1)
|
|
values = self.value_projection(values).view(B, S, H, -1)
|
|
|
|
out, attn = self.inner_attention(
|
|
queries,
|
|
keys,
|
|
values,
|
|
attn_mask
|
|
)
|
|
if self.mix:
|
|
out = ops.transpose(out, (0, 2, 1, 3)).contiguous()
|
|
out = out.view(B, L, -1)
|
|
|
|
return self.out_projection(out), attn |