《Attention is all you need》Pytorch实现

《Attention is all you need》Pytorch实现代码

• transformer_from_scratch.py

Self Attention

示意图

代码实现

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
				"""
				
				:param embed_size: int
				:param heads: int
				
				        """
				super(SelfAttention, self).__init__()
				
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (self.head_dim * heads == embed_size), "Embed size needs to be div by heads."

        # Embedding Layer
        self.values = nn.Linear(self.embed_size, self.head_dim, bias=False)
        self.keys = nn.Linear(self.embed_size, self.head_dim, bias=False)
        self.queries = nn.Linear(self.embed_size, self.head_dim, bias=False)

        # Out Layer
        self.fc_out = nn.Linear(self.head_dim * heads, embed_size)

    def forward(self, values, keys, queries, mask):
				"""
				
				:param values:  (N,value_len,heads,head_dim)
				:param keys:    (N,key_len,heads,head_dim)
				:param queries: (N,query_len,heads,head_dim)
				:param mask:    (N,heads,query_len,key_len)
				:return out:    (N,query_len,heads,head_dim)
				
				        """
				N = queries.shape[0]
        values_len, key_len, query_len = values.shape[1], keys.shape[1], queries.shape[1]

        # Split embedding into self.heads pieces
        values = values.reshape(N, values_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = queries.reshape(N, query_len, self.heads, self.head_dim)

        energy = torch.einsum("nqhd,nkhd->nhqk", queries, keys)

        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        # attention shape: (N,heads,query_len,key_len)
        attention = nn.Softmax(dim=3)(energy / (self.embed_size ** (1 / 2)))

        # value_len always equals key_len
        out = torch.einsum("nhql,nlhd->nqhd", attention, values).reshape(N,query_len,self.heads*self.head_dim)
        # out shape: (N,query_len,heads,head_dim)
        out = self.fc_out(out)

        return out

注意事项

energy = torch.einsum("nqhd,nkhd->nhqk", queries, keys)