text_classifier_tf2

text_classifier_tf2

多模型文本分类框架 支持TextCNN、BERT等

该开源项目提供基于TensorFlow 2的多模型文本分类框架。支持TextCNN、TextRNN、BERT等模型,集成词向量增强、对抗训练、对比学习等功能。框架适用于二分类和多分类任务,提供灵活配置选项。项目还包含交互式预测和批量测试工具,便于分析模型性能和错误案例。

文本分类深度学习模型训练方法评估指标模型部署Github开源项目

Text Classifier

Authour GitHub license

公众号文章:文本分类之Text-CNN/RNN/RCNN算法原理及工程实现
公众号文章:一篇文章带你走进词向量并掌握Word2Vec

此仓库是基于Tensorflow2.3的文本分类任务,通过直接配置可支持:

  • TextCNN/TextRNN/TextRCNN/Transformer/Bert/AlBert/DistilBert基本分类模型的训练
  • TextCNN/TextRNN/TextRCNN/Transformer的token可选用词粒度/字粒度
  • Word2Vec特征增强后接TextCNN/TextRNN/TextRCNN/Transformer
  • 支持Attention-TextCNN/TextRNN
  • FGM和PGD两种对抗方法的引入训练
  • 对比学习方法R-drop引入
  • 支持二分类和多分类,支持FocalLoss
  • 保存为pb文件可供部署
  • 项目代码支持交互测试和批量测试,批量测试可以观察错误分布和分析bad case

环境

  • python 3.7.10
  • tensorflow-gpu==2.3.0
  • tensorflow-addons==0.15.0
  • tqdm==4.50.2
  • gensim==3.8.3
  • jieba==0.42.1
  • pandas==1.1.3
  • scikit-learn==0.23.2
  • transformers==4.6.1
  • texttable==1.6.4

更新历史

日期版本描述
2018-12-01v1.0.0初始仓库
2020-10-20v2.0.0重构项目
2020-10-26v2.1.0加入F1、Precise、Recall分类指标,计算方式支持macro、micro、average、binary
2020-11-26v2.3.1加入focal loss用于改善标签分布不平衡的情况
2020-11-19v2.4.0增加每个类别的指标,重构指标计算逻辑
2021-03-02v2.5.0使用Dataset替换自己写的数据加载器来加载数据
2021-03-15v3.0.0支持仅使用TextCNN/TextRCNN进行数据训练(基于词粒度的token,使用随机生成的Embedding层)
2021-03-16v3.1.0支持取用Word2Vec的词向量后接TextCNN/TextRCNN进行数据训练;在log中打印配置
2021-03-17v3.1.1根据词频过滤一部分频率极低的词,不加入词表
2021-04-25v3.1.6通过配置可选GPU和CPU进行训练
2021-06-17v3.2.0增加字粒度的模型训练预测
2021-09-27v3.3.0增加测试集的批量测试
2021-11-01v4.0.0增加对抗训练,目前支持FGM和PGD两种方式;增加Bert微调分类训练;更换demo数据集
2021-11-24v4.2.0增加Transformer模型做文本分类、增加对比学习方法r-drop
2022-04-22v5.0.0批量测试打印bad_case以及预测混淆情况、文件夹检查、配置里面不再自己定义标签顺序、各类预训练模型支持

数据集

部分头条新闻数据集

使用

配置

在config.py中配置好各个参数,文件中有详细参数说明

训练分类器

配置好下列参数

classifier_config = {
    # 模型选择
    # 传统模型:TextCNN/TextRNN/TextRCNN/Transformer
    # 预训练模型:Bert/DistilBert/AlBert/RoBerta/Electra/XLNet
    'classifier': 'TextCNN',
    # 若选择Bert系列微调做分类,请在pretrained指定预训练模型的版本
    'pretrained': 'bert-base-chinese',
    # 训练数据集
    'train_file': 'data/train_dataset.csv',
    # 验证数据集
    'val_file': 'data/val_dataset.csv',
    # 测试数据集
    'test_file': 'data/test_dataset.csv',
    # 引入外部的词嵌入,可选word2vec、Bert
    # word2vec:使用word2vec词向量做特征增强
    # 不填写则随机初始化的Embedding
    'embedding_method': '',
    # token的粒度,token选择字粒度的时候,词嵌入(embedding_method)无效
    # 词粒度:'word'
    # 字粒度:'char'
    'token_level': 'word',
    # 去停用词,路径需要在上面的word2vec_config中配置,仅限非预训练微调使用
    'stop_words': True,
    # 是否去掉特殊字符
    'remove_special': True,
    # 不外接词嵌入的时候需要自定义的向量维度
    'embedding_dim': 300,
    # 存放词表的地方
    'token_file': 'data/word-token2id',
    # 类别列表
    'classes': ['家居', '时尚', '教育', '财经', '时政', '娱乐', '科技', '体育', '游戏', '房产'],
    # 模型保存的文件夹
    'checkpoints_dir': 'model/textcnn',
    # 模型保存的名字
    'checkpoint_name': 'textcnn',
    # 使用Textcnn模型时候设定卷集核的个数
    'num_filters': 64,
    # 学习率
    # 微调预训练模型时建议更小,设置5e-5
    'learning_rate': 0.0005,
    # 优化器选择
    # 可选:Adagrad/Adadelta/RMSprop/SGD/Adam/AdamW
    'optimizer': 'Adam',
    # 训练epoch
    'epoch': 100,
    # 最多保存max_to_keep个模型
    'max_to_keep': 1,
    # 每print_per_batch打印
    'print_per_batch': 100,
    # 是否提前结束
    'is_early_stop': True,
    # 是否引入attention
    # 注意:textrcnn不支持
    'use_attention': False,
    # attention大小
    'attention_size': 300,
    'patient': 8,
    'batch_size': 256,
    'max_sequence_length': 300,
    # 遗忘率
    'dropout_rate': 0.5,
    # 隐藏层维度
    # 使用textrcnn、textrnn和transformer中需要设定
    # 使用transformer建议设定为2048
    'hidden_dim': 256,
    # 编码器个数(使用transformer需要设定)
    'encoder_num': 1,
    # 多头注意力的个数(使用transformer需要设定)
    'head_num': 12,
    # 若为二分类则使用binary
    # 多分类使用micro或macro
    'metrics_average': 'micro',
    # 类别样本比例失衡的时候可以考虑使用
    'use_focal_loss': False,
    # 使用标签平滑
    # 主要用在预训练模型微调,直接训练小模型使用标签平滑会带来负面效果,慎用
    'use_label_smoothing': False,
    'smooth_factor': 0.1,
    # 是否使用GAN进行对抗训练
    'use_gan': False,
    # 目前支持FGM和PGD两种方法
    # fgm:Fast Gradient Method
    # pgd:Projected Gradient Descent
    'gan_method': 'pgd',
    # 对抗次数
    'attack_round': 3,
    # 使用对比学习,不推荐和对抗方法一起使用,效率慢收益不大
    'use_r_drop': False
}

配置完参数之后开始训练模型

# [train_classifier, interactive_predict, test, save_model, train_word2vec, train_sif_sentence_vec]
mode = 'train_classifier'
  • 训练结果

train_results_textcnn

测试

训练好模型直接可以开始测试,可以进行交互测试也可以批量测试

  • 交互测试
# [train_classifier, interactive_predict, test, save_model, train_word2vec, train_sif_sentence_vec]
mode = 'interactive_predict'  

交互测试结果
interactive_predict

  • 批量测试

在测试数据集配置上填和训练/验证集文件同构的文件地址

# 测试数据集
'test_file': 'data/test_dataset.csv',

模式设定为测试模式

# [train_classifier, interactive_predict, test, save_model, train_word2vec, train_sif_sentence_vec]
mode = 'test'

批量测试结果
batch_test

批量测试完会给出各个标签混淆的分布
error_dis

批量测试完同时会给出一份bad_case文件,方便更细致的纠正标签和判断模型效果
badcese

训练word2vec

在config.py中的mode中配置好词向量训练的相关参数,并在mode中选择train_word2vec并运行:

word2vec_config = {
    'stop_words': 'data/w2v_data/stop_words.txt',  # 停用词(可为空)
    'train_data': 'data/w2v_data/dataset.csv',  # 词向量训练用的数据
    'model_dir': 'model/word2vec_model',  # 词向量模型的保存文件夹
    'model_name': 'word2vec_model.pkl',  # 词向量模型名
    'word2vec_dim': 300,  # 词向量维度
    'min_count': 3,  # 最低保留词频大小
    # 选择skip-gram和cbow
    'sg': 'cbow'
}

# [train_classifier, interactive_predict, test, save_model, train_word2vec, train_sif_sentence_vec]
mode = 'train_word2vec'

公众号

相关问题欢迎在公众号反馈:

编辑推荐精选

iTerms

iTerms

企业专属的AI法律顾问

iTerms是法大大集团旗下法律子品牌,基于最先进的大语言模型(LLM)、专业的法律知识库和强大的智能体架构,帮助企业扫清合规障碍,筑牢风控防线,成为您企业专属的AI法律顾问。

SimilarWeb流量提升

SimilarWeb流量提升

稳定高效的流量提升解决方案,助力品牌曝光

稳定高效的流量提升解决方案,助力品牌曝光

Sora2视频免费生成

Sora2视频免费生成

最新版Sora2模型免费使用,一键生成无水印视频

最新版Sora2模型免费使用,一键生成无水印视频

Transly

Transly

实时语音翻译/同声传译工具

Transly是一个多场景的AI大语言模型驱动的同声传译、专业翻译助手,它拥有超精准的音频识别翻译能力,几乎零延迟的使用体验和支持多国语言可以让你带它走遍全球,无论你是留学生、商务人士、韩剧美剧爱好者,还是出国游玩、多国会议、跨国追星等等,都可以满足你所有需要同传的场景需求,线上线下通用,扫除语言障碍,让全世界的语言交流不再有国界。

讯飞绘文

讯飞绘文

选题、配图、成文,一站式创作,让内容运营更高效

讯飞绘文,一个AI集成平台,支持写作、选题、配图、排版和发布。高效生成适用于各类媒体的定制内容,加速品牌传播,提升内容营销效果。

热门AI辅助写作AI工具讯飞绘文内容运营AI创作个性化文章多平台分发AI助手
TRAE编程

TRAE编程

AI辅助编程,代码自动修复

Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。

AI工具TraeAI IDE协作生产力转型热门
商汤小浣熊

商汤小浣熊

最强AI数据分析助手

小浣熊家族Raccoon,您的AI智能助手,致力于通过先进的人工智能技术,为用户提供高效、便捷的智能服务。无论是日常咨询还是专业问题解答,小浣熊都能以快速、准确的响应满足您的需求,让您的生活更加智能便捷。

imini AI

imini AI

像人一样思考的AI智能体

imini 是一款超级AI智能体,能根据人类指令,自主思考、自主完成、并且交付结果的AI智能体。

Keevx

Keevx

AI数字人视频创作平台

Keevx 一款开箱即用的AI数字人视频创作平台,广泛适用于电商广告、企业培训与社媒宣传,让全球企业与个人创作者无需拍摄剪辑,就能快速生成多语言、高质量的专业视频。

即梦AI

即梦AI

一站式AI创作平台

提供 AI 驱动的图片、视频生成及数字人等功能,助力创意创作

下拉加载更多