GPT-2 Masked Self-Attention
GPT-2 Self-attention: 1- Creating queries, keys, and values
python
复制代码
batch_size = 1
seq_len = 10
embed_dim = 768
x = Tensor(np.random.randn(batch_size, seq_len, embed_dim), mindspore.float32)
from mindnlp._legacy.functional import split
from mindnlp.models.utils.utils import Conv1D
c_attn = Conv1D(3 * embed_dim, embed_dim)
query, key, value = split(c_attn(x), embed_dim, axis=2)
query.shape, key.shape, value.shape
def split_heads(tensor, num_heads, attn_head_size):
"""
Splits hidden_size dim into attn_head_size and num_heads
"""
new_shape = tensor.shape[:-1] + (num_heads, attn_head_size)
tensor = tensor.view(new_shape)
return ops.transpose(tensor, (0, 2, 1, 3)) # (batch, head, seq_length, head_features)
num_heads = 12
head_dim = embed_dim // num_heads
query = split_heads(query, num_heads, head_dim)
key = split_heads(key, num_heads, head_dim)
value = split_heads(value, num_heads, head_dim)
query.shape, key.shape, value.shape
GPT-2 Self-attention: 2- Scoring
python
复制代码
attn_weights = ops.matmul(query, key.swapaxes(-1, -2))
attn_weights.shape
max_positions = seq_len
bias = Tensor(np.tril(np.ones((max_positions, max_positions))).reshape(
(1, 1, max_positions, max_positions)), mindspore.bool_)
bias
python
复制代码
from mindnlp._legacy.functional import where, softmax
attn_weights = attn_weights / ops.sqrt(ops.scalar_to_tensor(value.shape[-1]))
query_length, key_length = query.shape[-2], key.shape[-2]
causal_mask = bias[:, :, key_length - query_length: key_length, :key_length].bool()
mask_value = Tensor(np.finfo(np.float32).min, dtype=attn_weights.dtype)
attn_weights = where(causal_mask, attn_weights, mask_value)
np.finfo(np.float32).min
attn_weights[0, 0]
attn_weights = softmax(attn_weights, axis=-1)
attn_weights.shape
attn_weights[0, 0]
attn_output = ops.matmul(attn_weights, value)
attn_output.shape
GPT-2 Self-attention: 3.5- Merge attention heads
python
复制代码
def merge_heads(tensor, num_heads, attn_head_size):
"""
Merges attn_head_size dim and num_attn_heads dim into hidden_size
"""
tensor = ops.transpose(tensor, (0, 2, 1, 3))
new_shape = tensor.shape[:-2] + (num_heads * attn_head_size,)
return tensor.view(new_shape)
attn_output = merge_heads(attn_output, num_heads, head_dim)
attn_output.shape
GPT-2 Self-attention: 4- Projecting
python
复制代码
c_proj = Conv1D(embed_dim, embed_dim)
attn_output = c_proj(attn_output)
attn_output.shape