基于tensorflow 1.x 的bert系列预训练模型工具

2021-10-14 10:11:00 浏览数 (1)

tfbert

  • 基于tensorflow 1.x 的bert系列预训练模型工具
  • 支持多GPU训练,支持梯度累积,支持pb模型导出,自动剔除adam参数
  • 采用dataset 和 string handle配合,可以灵活训练、验证、测试,在训练阶段也可以使用验证集测试模型,并根据验证结果保存参数。

源代码获取方式

关注微信公众号 datayx 然后回复 tf 即可获取。

AI项目体验地址 https://loveai.tech

说明

config、tokenizer参考的transformers的实现。

内置有自定义的Trainer,像pytorch一样使用tensorflow1.14,具体使用下边会介绍。

目前内置 文本分类、文本多标签分类、命名实体识别例子。

内置的几个例子的数据处理代码都支持多进程处理,实现方式参考的transformers。

内置代码示例数据集百度网盘提取码:rhxk

支持模型

bert、electra、albert、nezha、wobert、ChineseBert(GlyceBert)

requirements

代码语言:javascript复制
tensorflow==1.x
tqdm
jieba

目前本项目都是在tensorflow 1.x下实现并测试的,最好使用1.14及以上版本,因为内部tf导包都是用的

import tensorflow.compat.v1 as tf

使用说明

Config 和 Tokenizer

使用方法和transformers一样

多卡运行方式,需要设置环境变量CUDA_VISIBLE_DEVICES,内置trainer会读取参数:

代码语言:javascript复制
CUDA_VISIBLE_DEVICES=1,2 python run.py

详情查看代码样例

XLA和混合精度训练训练速度测试

使用哈工大的rbt3权重进行实验对比,数据为example中的文本分类数据集。开启xla和混合精度后刚开始训练需要等待一段时间优化,所以第一轮会比较慢, 等开启后训练速度会加快很多。最大输入长度32,批次大小32,训练3个epoch, 测试环境为tensorflow1.14,GPU是2080ti。

开启混合精度比较慢,base版本模型的话需要一两分钟,但是开启后越到后边越快,训练步数少的话可以只开启xla就行了,如果多的话 最好xla和混合精度(混合精度前提是你的卡支持fp16)都打开。

可加载中文权重链接

0 人点赞