def get_future_mask(self, batch_size, sequence_length):
"""Mask future targets and padding
:param batch_size: a Tensor dimension
:param sequence_length: a Tensor dimension
:param padding_mask: None or bool Tensor with shape [batch_size, sequence_length]
:return mask Tensor with shape [batch_size, sequence_length, sequence_length]
"""
xind = th.arange(sequence_length)[None,:].repeat(*(sequence_length, 1))
yind = th.arange(sequence_length)[:,None].repeat(*(1, sequence_length))
mask = yind >= xind
mask = mask[None,...].repeat(*(batch_size, 1, 1))
return mask.to(get_device())
def forward(self, inputs, start=1):
"""
Args:
inputs: a float32 Tensor with shape [batch_size, sequence_length, hidden_size]
Returns:
embedding: a float32 Tensor with shape [batch_size, sequence_length, hidden_size]
"""
#################################### YOUR CODE HERE ####################################
# PART 3: Implement the Position Embedding.
# As stated in section 3.5 of the paper, attention does not naturally embed position information
# To incorporate that, the authors use a variable frequency sin embedding.
# Note that we use zero-indexing here while the authors use one-indexing
assert inputs.shape[
-1] == self.hidden_size and 'Input final dim must match model hidden size'
batch_size = inputs.shape[0]
sequence_length = inputs.shape[1]
# obtain a sequence that starts at `start` and increments for `sequence_length `
seq_pos = th.arange(start, sequence_length + start, dtype=th.float32)
seq_pos_expanded = seq_pos[None, :, None]
index = seq_pos_expanded.repeat(*[1, 1, self.hidden_size // 2])
# create the position embedding as described in the paper
# use the `divisor` attribute instantiated in __init__
sin_embedding = th.sin(index / self.divisor)
cos_embedding = th.cos(index / self.divisor)
# interleave the sin and cos. For more info see:
# https://discuss.pytorch.org/t/how-to-interleave-two-tensors-along-certain-dimension/11332/3
position_shape = (1, sequence_length, self.hidden_size
) # fill in the other two dimensions
position_embedding = th.stack((sin_embedding, cos_embedding),
dim=3).view(position_shape)
pos_embed_deviced = position_embedding.to(get_device())
return inputs + position_embedding # add the embedding to the input
def forward(self, queries, keys, values, mask=None):
"""Fast scaled dot product attention.
:param queries: Tensor with shape [batch_size, heads (optional), n_queries, depth_k]
:param keys: Tensor with shape [batch_size, heads (optional), n_keyval, depth_k]
:param values: Tensor with shape [batch_size, heads (optional), n_keyval, depth_v]
:param mask: Tensor with shape [batch_size, n_queries, n_queries]
:return: output: Tensor with shape [batch_size, heads (optional), n_queries, depth_v]
"""
#################################### YOUR CODE HERE ####################################
# n_queries corresponds to the sequence length on the query side
# n_keyval corresponds to the sequence length on the key side (and value, as they are one and the same)
# depth_k is the size of the projection that the key / query comparison is performed on.
# depth_v is the size of the projection of the value projection. In a setting with one head, it is usually the dimension (dim) of the Transformer.
# heads corresponds to the number of heads the attention is performed on.
# If you are unfamiliar with attention heads, read section 3.2.2 of the Attention is all you need paper
# PART 1: Implement Attention QKV
# Use queries, keys and values to compute the output of the QKV attention
# As defined is the Attention is all you need paper: https://arxiv.org/pdf/1706.03762.pdf
device = get_device()
key_dim = th.tensor(keys.shape[-1], dtype=th.float32)
batch_size = queries.shape[0]
similarity = None
if len(queries.size()) == 3:
similarity = th.zeros(batch_size, queries.shape[1], keys.shape[1])
similarity = similarity.to(device)
for batch in range(batch_size):
similarity[batch, :, :] = th.mm(queries[batch, :, :],
keys[batch, :, :].t())
else:
heads = queries.shape[1]
similarity = th.zeros(batch_size, heads, queries.shape[2],
keys.shape[2])
similarity = similarity.to(device)
for batch in range(batch_size):
for head in range(heads):
similarity[batch,
head, :, :] = th.mm(queries[batch, head, :, :],
keys[batch, head, :, :].t())
similarity = (
1 / th.sqrt(key_dim)
) * similarity # Compute the similarity according to the QKV formula
masked_similarity = self.apply_mask(
similarity, mask=mask
) # We give you the mask to apply so that it is correct, you do not need to modify this.
last_dim = len(masked_similarity.size()) - 1
weights = F.softmax(
masked_similarity, dim=last_dim
) # Turn the similarity into a normalized output. Remember that the last dim contains the features
output = None
if len(queries.size()) == 3:
output = th.zeros(batch_size, queries.shape[1], values.shape[2])
output = output.to(device)
for batch in range(batch_size):
output[batch, :, :] = th.mm(weights[batch, :, :],
values[batch, :, :])
else:
heads = queries.shape[1]
output = th.zeros(batch_size, heads, queries.shape[2],
values.shape[3])
output = output.to(device)
for batch in range(batch_size):
for head in range(heads):
output[batch,
head, :, :] = th.mm(weights[batch, head, :, :],
values[batch, head, :, :])
#################################### END OF YOUR CODE ##################################
return output, weights