基于TensorFlow 2的多功能文本分类框架:text_classifier_tf2

Ray

文本分类是自然语言处理中的一项基础任务,在情感分析、主题分类、意图识别等多个领域有广泛应用。随着深度学习技术的发展,基于神经网络的文本分类模型在性能上已经远远超过了传统的机器学习方法。然而,要构建一个高质量的文本分类系统,不仅需要选择合适的模型,还需要考虑数据预处理、特征提取、模型训练等多个环节。为了简化这一过程,GitHub用户stanleylsx开发了text_classifier_tf2项目,旨在提供一个功能完备、易于使用的文本分类框架。

项目概述

text_classifier_tf2是一个基于TensorFlow 2.3的文本分类项目,支持多种主流分类模型和技巧。该项目的主要特点包括:

支持多种分类模型:TextCNN、TextRNN、TextRCNN、Transformer等传统模型,以及Bert、AlBert、DistilBert等预训练模型。
灵活的特征表示:支持词级和字级的token粒度,可选择使用Word2Vec或随机初始化的词嵌入。
丰富的训练技巧:支持注意力机制、对抗训练(FGM和PGD)、对比学习(R-drop)等方法提升模型性能。
完整的任务支持:支持二分类和多分类任务,提供focal loss等处理类别不平衡的方法。
便捷的使用方式:通过配置文件即可完成模型选择、参数设置等操作,无需修改代码。
全面的评估功能:支持交互式预测和批量测试,可输出详细的评估指标和错误分析。
部署支持:训练好的模型可保存为pb文件,方便在生产环境中部署。

环境配置

text_classifier_tf2基于Python 3.7和TensorFlow 2.3开发,主要依赖包括:

tensorflow-gpu==2.3.0
tensorflow-addons==0.15.0
gensim==3.8.3
jieba==0.42.1
pandas==1.1.3
scikit-learn==0.23.2
transformers==4.6.1

用户可以通过requirements.txt文件一键安装所需依赖。

使用方法

使用text_classifier_tf2进行文本分类主要包括以下步骤:

数据准备:将训练集、验证集和测试集按指定格式保存为CSV文件。
配置参数:在config.py文件中设置模型类型、数据路径、训练参数等。
训练模型:设置mode为'train_classifier'并运行main.py。
模型评估:设置mode为'test'进行批量测试,或设置为'interactive_predict'进行交互式预测。
模型部署:设置mode为'save_model'将模型保存为pb文件。

此外,项目还提供了训练Word2Vec词向量的功能,可以通过设置mode为'train_word2vec'来使用。

核心功能

多种分类模型

text_classifier_tf2支持以下几类文本分类模型:

传统模型:TextCNN、TextRNN、TextRCNN、Transformer
预训练模型:Bert、DistilBert、AlBert、RoBerta、Electra、XLNet

用户可以通过配置文件中的'classifier'参数选择使用的模型。对于预训练模型,还需要在'pretrained'参数中指定具体的版本。

特征表示

项目支持两种token粒度:

词级(word):使用jieba分词,可选择Word2Vec词向量或随机初始化的词嵌入
字级(char):直接使用单字作为token

通过'token_level'参数可以选择使用的粒度。当选择词级粒度时,可以通过'embedding_method'参数选择是否使用预训练的Word2Vec词向量。

训练技巧

为了提升模型性能,text_classifier_tf2实现了多种训练技巧:

注意力机制:通过'use_attention'参数开启,可用于TextCNN和TextRNN模型。
对抗训练:支持FGM(Fast Gradient Method)和PGD(Projected Gradient Descent)两种方法,通过'use_gan'和'gan_method'参数控制。
对比学习:实现了R-drop方法,通过'use_r_drop'参数开启。
标签平滑:主要用于预训练模型的微调,通过'use_label_smoothing'参数控制。
Focal Loss:用于处理类别不平衡问题,通过'use_focal_loss'参数开启。