pytorch基础 nn.embedding

nn.Embedding 是 PyTorch 中的一个模块，用于创建嵌入层（embedding layer），它将离散的索引（例如词汇表中的单词索引）映射为固定大小的稠密向量。这是许多 NLP 模型（包括 Transformer）中的基本组件。

示例用法：

import torch
import torch.nn as nn

# 定义一个嵌入层
vocab_size = 10000  # 词汇表大小
embedding_dim = 512  # 嵌入向量的维度
embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)

# 示例输入：一批词索引序列
# 形状：(batch_size, sequence_length)
input_indices = torch.tensor([[1, 2, 3], [4, 5, 6]])  # 2 个序列，每个序列长度为 3

# 获取嵌入向量
output_embeddings = embedding_layer(input_indices)
print(output_embeddings.shape)  # 输出形状：(batch_size, sequence_length, embedding_dim)

输出：

输出的张量形状为 (batch_size, sequence_length, embedding_dim)。对于上面的例子：

• batch_size = 2

• sequence_length = 3

• embedding_dim = 512

因此，输出形状为 (2, 3, 512)。

主要参数：

1. num_embeddings：词汇表的大小（唯一词的数量）, 通常比较大。

2. embedding_dim：嵌入向量的维度。

3. padding_idx：（可选）用于填充词的索引。如果提供，该索引对应的嵌入向量将固定为零向量，并且在训练过程中不会更新。

4. max_norm：（可选）如果指定，嵌入向量将被重新归一化，使其范数不超过该值。

5. norm_type：（可选）如果指定了 max_norm，则使用该范数类型（默认是 2，表示 L2 范数）。

6. scale_grad_by_freq：（可选）如果为 True，梯度将根据批次中词的频率进行缩放。

在 Transformer 中的使用：

在 Transformer 类中，nn.Embedding 通常用于创建：

1. 源嵌入：用于输入（源）序列。

2. 目标嵌入：用于输出（目标）序列。

3. 位置编码：为嵌入添加位置信息（可选，因为位置编码也可以单独添加）。

以下是在 Transformer 类中的使用示例：

class Transformer(nn.Module):
    def __init__(self, n_src_vocab, n_trg_vocab, src_pad_idx, trg_pad_idx, d_word_vec=512):
        super().__init__()
        
        # 源和目标嵌入层
        self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=src_pad_idx)
        self.trg_word_emb = nn.Embedding(n_trg_vocab, d_word_vec, padding_idx=trg_pad_idx)
        
        # Transformer 的其他组件...

这是对 nn.Embedding 及其在 Transformer 模型中的作用的简要介绍。如果你需要进一步的解释或示例，请告诉我！