nn/chapter10.ipynb

In [2]:

import torch
from d2l import torch as d2l

In [3]:

def show_heatmaps(matrices,xlabel,ylabel,titles=None,figsize=(2.5,2.5),cmap='Reds'):
    d2l.use_svg_display()
    num_rows,num_cols = matrices.shape[0],matrices.shape[1]
    fig,axes = d2l.plt.subplots(num_rows,num_cols,figsize=figsize,sharex=True,squeeze=False)
    for i,(row_axes,row_matrices) in enumerate(zip(axes,matrices)):
        for j,(ax,matrix) in enumerate(zip(row_axes,row_matrices)):
            pcm = ax.imshow(matrix.detach().numpy(),cmap=cmap)
            if i == num_rows - 1:
                ax.set_xlabel(xlabel)
            if j == 0:
                ax.set_ylabel(ylabel)
            if titles:
                ax.set_title(titles[j])
    fig.colorbar(pcm,ax=axes,shrink=0.6)

In [4]:

attention_weights = torch.eye(10).reshape((1, 1, 10, 10))

show_heatmaps(attention_weights, xlabel='Keys', ylabel='Queries')

In [5]:

n_train = 50
x_train,_ = torch.sort(torch.rand(n_train)*5)

In [6]:

def f(x):
    return 2*torch.sin(x)+x**0.8

In [7]:

y_train = f(x_train)+torch.normal(0.0,0.5,(n_train,))
x_test = torch.arange(0,5,0.1)
y_truth = f(x_test)
n_test = len(x_test)
n_test

Out [7]:

50

In [8]:

def plot_kernel_reg(y_hat):
    d2l.plot(x_test, [y_truth, y_hat], 'x', 'y', legend=['Truth', 'Pred'],
        xlim=[0, 5], ylim=[-1, 5])
    d2l.plt.plot(x_train, y_train, 'o', alpha=0.5);

In [9]:

y_hat = torch.repeat_interleave(y_train.mean(),n_test)
plot_kernel_reg(y_hat)

In [10]:

from torch import nn
X_repeat = x_test.repeat_interleave(n_train).reshape((-1, n_train))
attention_weights = nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim=1)
y_hat = torch.matmul(attention_weights,y_train)
x_train,X_repeat,attention_weights

Out [10]:

(tensor([9.4801e-04, 5.2621e-02, 2.3643e-01, 2.9541e-01, 3.2529e-01, 3.5294e-01,
         3.6399e-01, 4.1190e-01, 5.8751e-01, 6.8494e-01, 7.2698e-01, 8.0101e-01,
         9.2743e-01, 9.7225e-01, 1.1656e+00, 1.4397e+00, 1.4788e+00, 1.6088e+00,
         1.7147e+00, 1.8255e+00, 1.8951e+00, 1.9322e+00, 2.3849e+00, 2.4183e+00,
         2.4902e+00, 2.5441e+00, 2.6483e+00, 2.8371e+00, 2.9705e+00, 3.1383e+00,
         3.2709e+00, 3.3751e+00, 3.3849e+00, 3.4902e+00, 3.5679e+00, 3.5782e+00,
         3.5821e+00, 3.8648e+00, 4.1552e+00, 4.3914e+00, 4.4348e+00, 4.4414e+00,
         4.6038e+00, 4.6099e+00, 4.7481e+00, 4.7528e+00, 4.8005e+00, 4.8134e+00,
         4.8934e+00, 4.9044e+00]),
 tensor([[0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
         [0.1000, 0.1000, 0.1000,  ..., 0.1000, 0.1000, 0.1000],
         [0.2000, 0.2000, 0.2000,  ..., 0.2000, 0.2000, 0.2000],
         ...,
         [4.7000, 4.7000, 4.7000,  ..., 4.7000, 4.7000, 4.7000],
         [4.8000, 4.8000, 4.8000,  ..., 4.8000, 4.8000, 4.8000],
         [4.9000, 4.9000, 4.9000,  ..., 4.9000, 4.9000, 4.9000]]),
 tensor([[6.8662e-02, 6.8567e-02, 6.6769e-02,  ..., 6.3916e-07, 4.3363e-07,
          4.1091e-07],
         [6.4257e-02, 6.4501e-02, 6.3975e-02,  ..., 9.6787e-07, 6.6191e-07,
          6.2792e-07],
         [5.9935e-02, 6.0473e-02, 6.1093e-02,  ..., 1.4608e-06, 1.0070e-06,
          9.5634e-07],
         ...,
         [9.4878e-07, 1.2079e-06, 2.7906e-06,  ..., 5.8779e-02, 5.8062e-02,
          5.7935e-02],
         [6.1962e-07, 7.9294e-07, 1.8658e-06,  ..., 6.2113e-02, 6.1848e-02,
          6.1781e-02],
         [4.0301e-07, 5.1842e-07, 1.2425e-06,  ..., 6.5370e-02, 6.5614e-02,
          6.5615e-02]]))

In [11]:

plot_kernel_reg(y_hat)

In [12]:

d2l.show_heatmaps(attention_weights.unsqueeze(0).unsqueeze(0),
        xlabel='Sorted training inputs',
        ylabel='Sorted testing inputs')

In [13]:

class NWKernelRegression(nn.Module):
    def __init__(self,**kwargs):
        super().__init__(**kwargs)
        self.w = nn.Parameter(torch.randn((1,),requires_grad=True))
    def forward(self,queries,keys,values):
        queries = queries.repeat_interleave(keys.shape[1]).reshape(-1,keys.shape[1])
        self.attention_weights = nn.functional.softmax(
                -((queries - keys) * self.w)**2 / 2, dim=1)
        return torch.bmm(self.attention_weights.unsqueeze(1),values.unsqueeze(-1)).reshape(-1)

In [14]:

# X_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输入
X_tile = x_train.repeat((n_train, 1))
# Y_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输出
Y_tile = y_train.repeat((n_train, 1))
# keys的形状:('n_train'，'n_train'-1)
keys = X_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))
# values的形状:('n_train'，'n_train'-1)
values = Y_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))

In [15]:

net = NWKernelRegression()
loss = nn.MSELoss(reduction='none')
trainer = torch.optim.SGD(net.parameters(), lr=0.5)
animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[1, 5])
for epoch in range(5):
    trainer.zero_grad()
    l = loss(net(x_train, keys, values), y_train)
    l.sum().backward()
    trainer.step()
    print(f'epoch {epoch + 1}, loss {float(l.sum()):.6f}')
    animator.add(epoch + 1, float(l.sum()))

In [16]:

# keys的形状:(n_test，n_train)，每一行包含着相同的训练输入（例如，相同的键）
keys = x_train.repeat((n_test, 1))
# value的形状:(n_test，n_train)
values = y_train.repeat((n_test, 1))
y_hat = net(x_test, keys, values).unsqueeze(1).detach()
plot_kernel_reg(y_hat)

In [17]:

d2l.show_heatmaps(net.attention_weights.unsqueeze(0).unsqueeze(0),
xlabel='Sorted training inputs',
ylabel='Sorted testing inputs')

In [18]:

def masked_softmax(X,valid_lens):
    # X:3D valid_len 1D or 2D
    if valid_lens is None:
        return nn.functional.softmax(X,dim=-1)
    else:
        shape = X.shape
        if valid_lens.dim()==1:
            valid_lens = torch.repeat_interleave(valid_lens,shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        #print(valid_lens)
        X = d2l.sequence_mask(X.reshape(-1,shape[-1]),valid_lens,value=-1e6)
        return nn.functional.softmax(X.reshape(shape),dim=-1)

In [19]:

masked_softmax(torch.rand(2,2,4),torch.tensor([2,3]))

Out [19]:

tensor([[[0.6937, 0.3063, 0.0000, 0.0000],
         [0.6165, 0.3835, 0.0000, 0.0000]],

        [[0.3526, 0.3291, 0.3183, 0.0000],
         [0.2735, 0.2359, 0.4906, 0.0000]]])

In [20]:

class AdditiveAttention(nn.Module):
    def __init__(self,key_size,query_size,num_hiddens,dropout,**kwargs):
        super(AdditiveAttention,self).__init__(**kwargs)
        self.W_k = nn.Linear(key_size,num_hiddens,bias=False)
        self.W_q = nn.Linear(query_size,num_hiddens,bias=False)
        self.w_v = nn.Linear(num_hiddens,1,bias=False)
        self.dropout=nn.Dropout(dropout)
    def forward(self,queries,keys,value,valid_lens):
        queries,keys=self.W_q(queries),self.W_k(keys)
        # queries (batch_size,n_q,1,num_hidden)
        # key     (batch_size,1,n_k,num_hiddens)
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        #features (batch_size,n_q,n_k,num_hidden)
        features = torch.tanh(features)
        scores = self.w_v(features).squeeze(-1)
        #print(f"Inside AdditiveAttention: value.shape = {value.shape}")  # 检查此处形状
        self.attention_weights = masked_softmax(scores, valid_lens)

        return torch.bmm(self.dropout(self.attention_weights), value)

In [21]:

queries,keys = torch.normal(0,1,(2,1,20)) , torch.ones((2,10,2))
values =  torch.arange(40,dtype=torch.float32).reshape(1,10,4).repeat(2,1,1)
valid_lens = torch.tensor([2, 6])
attention = AdditiveAttention(key_size=2, query_size=20, num_hiddens=8,
                                dropout=0.1)
attention.eval()
attention(queries,keys,values,valid_lens)

Out [21]:

tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],

        [[10.0000, 11.0000, 12.0000, 13.0000]]], grad_fn=<BmmBackward0>)

In [22]:

d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
        xlabel='Keys', ylabel='Queries')

In [23]:

import math
class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
# queries的形状：(batch_size，查询的个数，d)
# keys的形状：(batch_size，“键－值”对的个数，d)
# values的形状：(batch_size，“键－值”对的个数，值的维度)
# valid_lens的形状:(batch_size，)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # 设置transpose_b=True为了交换keys的最后两个维度
        scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)

        return torch.bmm(self.dropout(self.attention_weights), values)

In [24]:

queries = torch.normal(0, 1, (2, 1, 2))
attention = DotProductAttention(dropout=0.5)
attention.eval()
attention(queries, keys, values, valid_lens)

Out [24]:

tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],

        [[10.0000, 11.0000, 12.0000, 13.0000]]])

In [25]:

d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
xlabel='Keys', ylabel='Queries')

In [26]:

class AttentionDecoder(d2l.Decoder):
    def __init__(self,**kwargs):
        super(AttentionDecoder, self).__init__(**kwargs)
    @property
    def attention_weight(self):
        raise NotImplementedError

In [27]:

class Seq2SeqAttentionDecoder(AttentionDecoder):
    def __init__(self,vocab_size,embed_size,num_hiddens,num_layers,dropout=0,**kwargs):
        super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)
        self.attention = AdditiveAttention(
        num_hiddens,  num_hiddens, num_hiddens,dropout)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(
        embed_size + num_hiddens, num_hiddens, num_layers,
        dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)
    def init_state(self, enc_outputs, enc_valid_lens, *args):
# outputs的形状为(batch_size，num_steps，num_hiddens).
# hidden_state的形状为(num_layers，batch_size，num_hiddens)
        outputs, hidden_state = enc_outputs
        #print(f"Encoder outputs shape before permute: {outputs.shape}")  # 应为 (num_steps, batch_size, num_hiddens) 或 (batch_size, num_steps, num_hiddens)
        enc_outputs_permuted = outputs.permute(1, 0, 2)
        #print(f"After permute: {enc_outputs_permuted.shape}")  # 期望 (batch_size, num_steps, num_hiddens)
        return (enc_outputs_permuted, hidden_state, enc_valid_lens)
    def forward(self, X, state):
        # enc_outputs的形状为(batch_size,num_steps,num_hiddens).
        # hidden_state的形状为(num_layers,batch_size,
        # num_hiddens)
        enc_outputs, hidden_state, enc_valid_lens = state
        # 输出X的形状为(num_steps,batch_size,embed_size)
        X = self.embedding(X).permute(1, 0, 2)
        outputs, self._attention_weights = [], []
        for x in X:
            # query的形状为(batch_size,1,num_hiddens)
            query = torch.unsqueeze(hidden_state[-1], dim=1)
            # context的形状为(batch_size,1,num_hiddens)
            #print(f"values shape before attention: {enc_outputs.shape}")  # 应为 (4, 7, 16)
            context = self.attention(
            query, enc_outputs, enc_outputs, enc_valid_lens)
            # 在特征维度上连结
            x = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)
            # 将x变形为(1,batch_size,embed_size+num_hiddens)
            out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)
            outputs.append(out)
            self._attention_weights.append(self.attention.attention_weights)
        # 全连接层变换后，outputs的形状为
        # (num_steps,batch_size,vocab_size)
        outputs = self.dense(torch.cat(outputs, dim=0))
        return outputs.permute(1, 0, 2), [enc_outputs, hidden_state,
        enc_valid_lens]
    @property
    def attention_weights(self):
        return self._attention_weights

In [28]:

encoder = d2l.Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,
num_layers=2)
encoder.eval()
decoder = Seq2SeqAttentionDecoder(vocab_size=10, embed_size=8, num_hiddens=16,
num_layers=2)
decoder.eval()
X = torch.zeros((4, 7), dtype=torch.long) # (batch_size,num_steps)
state = decoder.init_state(encoder(X), None)
output, state = decoder(X, state)
output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape

Out [28]:

(torch.Size([4, 7, 10]), 3, torch.Size([4, 7, 16]), 2, torch.Size([4, 16]))

In [28]:

In [29]:

embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1
batch_size, num_steps = 64, 10
lr, num_epochs, device = 0.005, 250, d2l.try_gpu()
train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)
encoder = d2l.Seq2SeqEncoder(
len(src_vocab), embed_size, num_hiddens, num_layers, dropout)
decoder = Seq2SeqAttentionDecoder(
len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)
net = d2l.EncoderDecoder(encoder, decoder)
d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

loss 0.019, 17862.0 tokens/sec on cuda:0

In [30]:

engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
                f'bleu {d2l.bleu(translation, fra, k=2):.3f}')

go . => va !,  bleu 1.000
i lost . => j'ai perdu .,  bleu 1.000
he's calm . => il est <unk> .,  bleu 0.658
i'm home . => je suis chez moi .,  bleu 1.000

In [32]:

class MultiHeadAttention(nn.Module):
    def __init__(self, key_size, query_size, value_size, num_hiddens,
             num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)
    def forward(self, queries, keys, values, valid_lens):
        # 1. 线性投影 + 变换形状以分割多头
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys    = transpose_qkv(self.W_k(keys),    self.num_heads)
        values  = transpose_qkv(self.W_v(values),  self.num_heads)

        # 2. 处理有效长度掩码（valid_lens）以适配多头
        if valid_lens is not None:
            valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)

        # 3. 计算注意力（每个头独立计算）
        output = self.attention(queries, keys, values, valid_lens)

        # 4. 合并多头，并通过输出线性层
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)
def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，
    # num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)
    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)
    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])

def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

In [33]:

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
num_hiddens, num_heads, 0.5)
attention.eval()

Out [33]:

MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)

In [34]:

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens = 6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape

Out [34]:

torch.Size([2, 4, 100])

In [ ]: