【pytorch】nn.Embedding()

### PyTorch `nn.Embedding` 的用法 在自然语言处理 (NLP) 和其他领域中，嵌入层 (`Embedding Layer`) 是非常重要的组件之一。PyTorch 提供了 `torch.nn.Embedding` 类来创建和管理这些嵌入。 #### 创建 Embedding 层 可以通过指定词汇表大小（即词典中的单词数量）以及每个词的维度来初始化一个 embedding 层： ```python import torch from torch import nn embedding = nn.Embedding(num_embeddings=10, embedding_dim=3) ``` 这里 `num_embeddings` 表示词汇量大小，而 `embedding_dim` 则表示每个词语对应的向量长度[^1]。 #### 输入与输出形状 对于输入张量而言，其最后一个维度应当代表 batch 中样本的数量；其余所有前置维度都将被视作单一样本的一部分。因此，在大多数情况下，输入会是一个二维张量 `(batch_size, seq_len)` 或者一维张量 `(seq_len,)` 当批次大小为 1 时。该函数将会返回具有相同前导尺寸的新张量，但在最后增加了一个新的特征轴用于存储 embeddings 结果，形成三维张量 `(batch_size, seq_len, embedding_dim)` 或两维张量 `(seq_len, embedding_dim)` 如果批次大小为 1。 #### 使用预训练权重 如果希望利用已经预先训练好的 word vectors 初始化 embedding 层，则可以这样做: ```python pretrained_weights = ... # 加载预训练权重矩阵 vocab_size, embed_size = pretrained_weights.shape embedding_layer = nn.Embedding(vocab_size, embed_size).from_pretrained(pretrained_weights) ``` 注意这里的 `.from_pretrained()` 方法可以直接接受 numpy 数组或其他形式的数据作为参数，并据此设置初始权重值。 #### 实际应用案例 考虑如下场景：给定一批句子 ID 序列，想要获取它们各自的 embedding 向量以便后续处理。假设我们有一个简单的例子，其中包含三个句子，每句话有两个 token IDs 组成的小序列： ```python input_tensor = torch.LongTensor([[1, 2], [4, 5], [7, 8]]) output_embedding = embedding(input_tensor) print(output_embedding.size()) # 输出应形如: torch.Size([3, 2, 3]) ``` 这段代码片段展示了如何将整数型 tensor 转换成相应的浮点型 embedding vector 形式的操作过程。