lite-transformer

Lite Transformer 项目介绍

项目概述

Lite Transformer 是一种先进的神经网络模型，主要用于自然语言处理任务。它通过引入长短程注意力机制，对传统 Transformer 模型进行了优化。此模型通过减少计算量并提高效率，实现在多个重要数据集上的高效表现。

如何使用

环境依赖

要使用 Lite Transformer，用户需要以下软件环境：

• Python 版本需要在 3.6 以上
• PyTorch 版本需在 1.0.0 以上
• 需要 configargparse 包，至少为 0.14 版本
• 若计划训练新的模型，还需 NVIDIA GPU 以及 NCCL 库的支持

安装步骤

1. 代码库安装

   用户可以通过 pip 从源码安装 fairseq 来本地开发：

   pip install --editable .

2. 自定义模块

   为了支持 GPU，还需要构建 lightconv 和 dynamicconv 模块。

   • Lightconv_layer 模块：

     cd fairseq/modules/lightconv_layer
     python cuda_function_gen.py
     python setup.py install

   • Dynamicconv_layer 模块：

     cd fairseq/modules/dynamicconv_layer
     python cuda_function_gen.py
     python setup.py install

数据准备

Lite Transformer 需要对数据进行预处理，支持多个数据集：

1. IWSLT'14 De-En
2. WMT'14 En-Fr
3. WMT'16 En-De
4. WIKITEXT-103

使用 bash 脚本来下载和预处理这些数据集。例如：

bash configs/iwslt14.de-en/prepare.sh

测试步骤

要在 WMT'14 En-Fr 数据集上测试模型，可以执行以下命令：

configs/wmt14.en-fr/test.sh [model checkpoin路径] [gpu-id] [test|valid]

例如，在 GPU 0 上评估 Lite Transformer：

configs/wmt14.en-fr/test.sh embed496/ 0 test

模型训练

Lite Transformer 提供了一些示例，帮助用户训练模型。以在 WMT'14 En-Fr 数据集上使用 8 个 GPU 训练为例：

python train.py data/binary/wmt14_en_fr --configs configs/wmt14.en-fr/attention/multibranch_v2/embed496.yml

如果 GPU 数量少，例如 4 个，可以调整为：

CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py data/binary/wmt14_en_fr --configs configs/wmt14.en-fr/attention/multibranch_v2/embed496.yml --update-freq 32

分布式训练（可选）

Lite Transformer 支持在多节点上进行分布式训练。例如，使用两个节点共 16 个 GPU：

• 在 Host1 上：

  python -m torch.distributed.launch \
          --nproc_per_node=8 \
          --nnodes=2 --node_rank=0 \
          --master_addr=host1 --master_port=8080 \
          train.py data/binary/wmt14_en_fr \
          --configs configs/wmt14.en-fr/attention/multibranch_v2/embed496.yml \
          --distributed-no-spawn \
          --update-freq 8

• 在 Host2 上：

  python -m torch.distributed.launch \
          --nproc_per_node=8 \
          --nnodes=2 --node_rank=1 \
          --master_addr=host1 --master_port=8080 \
          train.py data/binary/wmt14_en_fr \
          --configs configs/wmt14.en-fr/attention/multibranch_v2/embed496.yml \
          --distributed-no-spawn \
          --update-freq 8

模型成果

Lite Transformer 提供了多个数据集的模型检查点，这些数据集包括 WMT'14 En-Fr、WMT'16 En-De、CNN/DailyMail 和 WIKITEXT-103 等。这些模型在测试集上的表现如表所示，用户可以直接下载并使用这些已训练的模型。

表中列出的模型指标与数据集的下载链接，方便用户获取并利用这些高性能的预训练模型进行研究和开发。