April 23, 2020

更亲开发者的深度学习包: PYTORCH LIGHTNING 2 -- 在 TPU 上用6分6秒跑完 GermEval 2014 NER 任务的 3 个 Epoch

更亲开发者的深度学习包: PYTORCH LIGHTNING 2 -- 在 TPU 上用6分6秒跑完 GermEval 2014 NER 任务的 3 个 Epoch

拿 transformers 里面的 NER 例子在 Colab TPU 环境下跑一下,这个例子是 GermEval 2014 NER 任务,demo 中默认 Epoch 设为 3,跑完后 dev 上的 F1 可以到 84%。下面按照 notebook 里面介绍走一遍。

首先介绍了一个 colab 上的 trick,上来就把内存占满造成 vm 崩溃的方式获得重启后挂在更多的内存,鉴于选择 TPU 之后 RAM 默认已经有 35G(gpu 的话有 25G),感觉已经够用了,而且现在执行这个小 trick 之后好像也没什么卵用了,所以就简单提一下就算了。

import torch; torch.tensor([10.]*10000000000)

开始肯定还是在 colab 上把 torch-xla 配置好,当然前提是你在笔记本设置里面选定了 TPU先, 加上 %%capture 是为了让 notebook 不要输出安装信息了,在本例中只要这部分不出现异常我们不关心执行细节。至于 xla 的版本,此时此刻选择 20200325 就好:

%%capture

import os
print(os.environ["COLAB_TPU_ADDR"])

assert os.environ["COLAB_TPU_ADDR"], "Make sure to select TPU from Edit > Notebook settings > Hardware accelerator"

VERSION = "20200325"  #@param ["1.5" , "20200325", "nightly"]
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version $VERSION

下面配置特定的 pytorch-lightning 和 transformers 版本,推荐采用这里用的 clip 版本,master 版本目前对 xla 的支持其实不稳定。除了这点之外,安装配置 transformers 也需要注意把所有的 requirements 安装好,并启用对应的 branch (tpu_pl2),这几个设置很重要,否则走到具体的 training 阶段就会各种卡壳

%%capture

!pip install git+git://github.com/srush/pytorch-lightning.git --upgrade
!git clone https://github.com/srush/transformers.git
!cd transformers; pip install -e . ; pip install -r examples/requirements.tx

数据清洗好:

%%bash
cd transformers/examples/ner/
curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-train.tsv?attredirects=0&d=1' \
| grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > train.txt.tmp
curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-dev.tsv?attredirects=0&d=1' \
| grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > dev.txt.tmp
curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-test.tsv?attredirects=0&d=1' \
| grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > test.txt.tmp
wget "https://raw.githubusercontent.com/stefan-it/fine-tuned-berts-seq/master/scripts/preprocess.py"
export MAX_LENGTH=128
export BERT_MODEL=bert-base-multilingual-cased
python3 preprocess.py train.txt.tmp $BERT_MODEL $MAX_LENGTH > train.txt
python3 preprocess.py dev.txt.tmp $BERT_MODEL $MAX_LENGTH > dev.txt
python3 preprocess.py test.txt.tmp $BERT_MODEL $MAX_LENGTH > test.txt
cat train.txt dev.txt test.txt | cut -d " " -f 2 | grep -v "^$"| sort | uniq > labels.tx

把 transformers 里面带的 ner 例子目录添加到 python 搜索路径中:

# Filter out warnings.
%tensorflow_version 2.x
import warnings
warnings.filterwarnings("ignore", message=".*add_.*")
import argparse, sys, os, shlex
sys.path.append("transformers/examples/ner/"

构造好 Trainer 并执行, 添加 %%time  是为了统计出来总共执行的时间:

%%time

from run_pl_ner import NERTransformer
from transformer_base import generic_train, add_generic_args

argString = "--data_dir transformers/examples/ner/ --model_type bert --labels transformers/examples/ner/labels.txt" + \
" --model_name_or_path bert-base-multilingual-cased --output_dir germ --max_seq_length 128 --num_train_epochs 3 " + \
"--train_batch_size 32 --eval_batch_size 32 --do_train --n_tpu_cores 8"

parser = argparse.ArgumentParser()                                                                                                                                                                               
add_generic_args(parser, os.getcwd())                                                                                                                                                                            
parser = NERTransformer.add_model_specific_args(parser, os.getcwd())                                                                                                                                             
args = parser.parse_args(shlex.split(argString))
model = NERTransformer(args)                                                                                              
trainer = generic_train(model, args) 

3 个 epoch 返回的 Evaluation 信息和最终信息如下:

这个例子说明了如何在 colab tpu 上面跑 transformers , 接下来研究下如何把 transformers 的这些场景用 PyTorch Lightning 改写。