如何在Google Colab上用TPU通过TensorFlow运行SkipGram版Word2Vec

Google Colab支持使用TPU运行TensorFlow实现的SkipGram模式Word2Vec模型，相比CPU/GPU训练速度可提升5~10倍，具体实现步骤如下：

• 打开Colab笔记本，点击菜单栏「修改」→「笔记本设置」，在「硬件加速器」选项中选择「TPU」，保存设置
• 导入依赖后完成TPU初始化，对应代码如下：

import tensorflow as tf

# 初始化TPU集群
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
# 定义TPU分布式策略
strategy = tf.distribute.TPUStrategy(resolver)

TPU要求所有数据处理逻辑基于TensorFlow原生算子实现，避免使用非TensorFlow实现的Python自定义逻辑，SkipGram的处理流程需要符合以下规范：

• 提前构建完成词汇表、词-ID映射表，转成TensorFlow静态常量存储
• 上下文滑动窗口构建、负采样逻辑直接用tf.data.Dataset原生接口实现，避免训练过程中CPU与TPU频繁交互产生的传输开销
• 批量大小建议设置为2的幂次，比如1024、2048，适配TPU计算单元的对齐要求，提升计算效率
• 数据流水线最后添加prefetch(tf.data.AUTOTUNE)配置，实现数据预加载，避免TPU空闲等待数据

所有模型变量、优化器、损失函数都要放在strategy.scope()上下文管理器中定义，才能正确分发到TPU的8个核心上并行计算，对应代码如下：

# 以下参数请根据自己的数据集调整
VOCAB_SIZE = 你的词汇表总大小
EMBEDDING_DIM = 128 # 词向量维度
BATCH_SIZE = 2048 # 全局批量大小

with strategy.scope():
    # SkipGram模型定义
    class SkipGram(tf.keras.Model):
        def __init__(self, vocab_size, embed_dim):
            super().__init__()
            self.target_embedding = tf.keras.layers.Embedding(
                input_dim=vocab_size,
                output_dim=embed_dim,
                input_length=1,
                name="w2v_embedding"
            )
            self.context_embedding = tf.keras.layers.Embedding(
                input_dim=vocab_size,
                output_dim=embed_dim,
                input_length=1
            )
        
        def call(self, pair):
            target, context = pair
            target_emb = self.target_embedding(target)
            context_emb = self.context_embedding(context)
            # 计算点积相似度
            dots = tf.einsum('be,bce->bc', target_emb, context_emb)
            return dots
    
    # 实例化模型、定义损失函数和优化器
    model = SkipGram(VOCAB_SIZE, EMBEDDING_DIM)
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
    # 适配分布式训练的损失计算
    def distributed_loss(labels, logits):
        per_example_loss = loss_fn(labels, logits)
        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=BATCH_SIZE)
    
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

# 编译模型
model.compile(optimizer=optimizer, loss=distributed_loss, metrics=['accuracy'])
# 启动训练，train_dataset为你提前构建好的tf.data格式训练数据集
model.fit(train_dataset, epochs=20)

训练完成后直接从Embedding层取出权重即可得到训练好的词向量：

word_vectors = model.get_layer('w2v_embedding').get_weights()[0]

• 训练过程中不要频繁打印日志、保存中间权重，会打断TPU的计算流水线，大幅拉低训练效率
• 负采样数量建议设置为5~20，在训练效果和速度上取平衡
• 数据集较大时建议先转成TFRecord格式存储，进一步提升数据加载速度

内容的提问来源于stack exchange，提问作者krcoder