nn.Embedding使用示例

假设有一个文本数据集，词汇表大小为10000个词，我们想要将每个词映射到一个128维的向量空间。

import torch  
import torch.nn as nn  
  
# 创建嵌入层实例  
embedding_layer = nn.Embedding(  
    num_embeddings=10000,  # 词汇表大小：10000个词  
    embedding_dim=128,    # 每个词的嵌入维度：128维  
    padding_idx=0         # 将索引0用作填充标记，其嵌入向量将全为零且在训练时不更新  
)  
  
# 假设有一个批次的样本，其中包含4个句子，每个句子被表示为一个词索引序列  
# 这里使用随机生成的索引作为示例  
batch_size = 4  
seq_len = 10  
sample_indices = torch.randint(1, 10000, (batch_size, seq_len))  # 生成随机索引，范围从1到9999（不包含0，因为0是填充索引）  
  
# 通过嵌入层获取嵌入向量  
embedded_samples = embedding_layer(sample_indices)  
  
# embedded_samples 的形状是 [batch_size, seq_len, embedding_dim]  
# [4, 10, 128]  
print(embedded_samples.shape)  # 输出：torch.Size([4, 10, 128])

上面代码中，nn.Embedding层将每个词索引映射到了一个128维的向量。sample_indices是一个形状为[batch_size, seq_len]的张量，表示一个批次中的句子，其中每个句子由词索引组成。通过将这些索引传递给embedding_layer，我们得到了一个形状为[batch_size, seq_len, embedding_dim]的张量，其中包含了每个词的嵌入向量。