### 如何在 PyTorch 中实现注意力机制 #### 使用自定义层构建基本注意力模块 为了实现在深度学习模型中的注意力机制，通常会创建一个专门用于计算注意力权重的函数或类。下面是一个简单的例子，在这个例子中实现了加法型(Additive)注意力机制： ```python import torch.nn as nn import torch class AdditiveAttention(nn.Module): def __init__(self, query_size, key_size, num_hiddens, dropout=0., **kwargs): super().__init__() 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, values, valid_lens=None): """queries: (batch_size, seq_len_queries, query_size), keys: (batch_size, seq_len_keys, key_size), values: (batch_size, seq_len_keys, value_size), valid_lens: either (batch_size,) or (batch_size, seq_len_queries).""" # After applying the linear transformation to `keys` and `queries`, they are added together. features = torch.tanh(self.W_k(keys.unsqueeze(1)) + self.W_q(queries.unsqueeze(2))) scores = self.w_v(features).squeeze(-1) # Shape of `scores`: (`batch_size`, `seq_len_queries`, `seq_len_keys`) attention_weights = masked_softmax(scores, valid_lens) outputs = torch.bmm(attention_weights, values) return self.dropout(outputs), attention_weights def masked_softmax(X, valid_lens): """ Perform softmax operation by masking elements on the last axis. X: 3D tensor with shape (`batch_size`, ..., `num_features`) valid_lens: 1D or 2D tensor indicating valid lengths for each item in batch, where None means no mask is applied. Returns a tensor after performing softmax over the specified dimension while ignoring padded positions.""" if valid_lens is None: return F.softmax(X, dim=-1) else: shape = X.shape if valid_lens.dim() == 1: valid_lens = torch.repeat_interleave(valid_lens, repeats=shape[1], dim=0) else: T = torch.tril(torch.ones((shape[1], shape[1]), dtype=torch.float32)) valid_lens = valid_lens[:, ::-1].cumsum(dim=1)[:, ::-1] # On the final axis, replace masked elements with a very small negative number before computing exp(x). X_exp = torch.exp(X.reshape(shape[0]*shape[1], -1)).reshape(*X.shape) indices = [slice(None)] * len(X.shape[:-1]) + [valid_lens.long()] X_exp_clipped = X_exp.clone() X_exp_clipped.scatter_(-1, indices, float('-inf')) return F.softmax(X_exp_clipped, dim=-1) ``` 上述代码展示了如何通过线性变换来组合查询向量(`query`)和键向量(`key`), 并利用另一个全连接层生成得分矩阵(score matrix)，最后经过softmax操作得到注意力分布并据此对值向量(`value`)加权求和[^1]。 对于更复杂的场景如自然语言处理任务，则可能会涉及到多头注意力(Multi-head Attention)结构以及Transformer架构下的编码器解码器框架内的具体实现方式；而对于计算机视觉领域特别是图像分类、目标检测等问题来说，除了全局平均池化外还可以考虑引入局部区域特征增强方案比如Squeeze-and-Excitation Networks(SENet)[^2] 或者文中提到的大选择性核网络(LSKNet)等方法进一步提升性能表现。