Pytorch TextCNN实现中文文本分类(附完整训练代码)_pytorch 中文文本分类

阿里云国内75折 回扣 微信号:monov8
阿里云国际,腾讯云国际,低至75折。AWS 93折 免费开户实名账号 代冲值 优惠多多 微信号:monov8 飞机:@monov6

Pytorch TextCNN实现中文文本分类附完整训练代码

目录

Pytorch TextCNN实现中文文本分类(附完整训练代码)

一、项目介绍

二、中文文本数据集

1THUCNews文本数据集

2 今日头条文本数据集 

3自定义文本数据集

三、TextCNN模型结构

1TextCNN模型结构

2TextCNN实现

四、训练词嵌入word2vec可选

五、文本预处理

1句子分词处理jieba中文分词

2特殊字符处理

3文本数据增强

六、训练过程 

1项目框架说明

2准备Train和Test文本数据

3配置文件config_textfolder.yaml

4开始训练

5可视化训练过程

6一些优化建议

七. 模型测试效果

八.项目源码下载


一、项目介绍

本篇将分享一个NLP项目实例利用深度学习框架Pytorch构建TextCNN模型也支持TextCNN,LSTM,BiLSTM模型实现一个简易的中文文本分类模型基于该项目训练的TextCNN的文本分类模型在THUCNews数据集上训练集的Accuracy 99%左右测试集的Accuracy在88.36%左右。

如果你想学习中文单词预测请参考《Pytorch LSTM实现中文单词预测(附完整训练代码)_

【尊重原则转载请注明出处】https://blog.csdn.net/guyuealian/article/details/127846717

二、中文文本数据集

中文文本数据集特别多这里仅仅介绍2个常用的文本文本分类数据集

1THUCNews文本数据集

THUCNews是根据新浪新闻RSS订阅频道2005~2011年间的历史数据筛选过滤生成包含74万篇新闻文档2.19 GB均为UTF-8纯文本格式。我们在原始新浪新闻分类体系的基础上重新整合划分出14个候选分类类别财经、彩票、房产、股票、家居、教育、科技、社会、时尚、时政、体育、星座、游戏、娱乐。使用THUCTC工具包在此数据集上进行评测准确率可以达到88.6%。

  • 官方数据集下载链接: http://thuctc.thunlp.org/message
  • 百度网盘下载链接: https://pan.baidu.com/s/1DT5xY9m2yfu1YGaGxpWiBQ 提取码: bbpe
  • THUCTC: 一个高效的中文文本分类工具包: THUCTC: 一个高效的中文文本分类工具

今日头条文本数据集 

今日头条文本数据集数据来源于今日头条客户端约382688条分布于15个分类中。

数据格式

6552431613437805063_!_102_!_news_entertainment_!_谢娜为李浩菲澄清网络谣言之后她的两个行为给自己加分_!_佟丽娅,网络谣言,快乐大本营,李浩菲,谢娜,观众们

每行为一条数据以_!_分割的个字段从前往后分别是 新闻ID分类code见下文分类名称见下文新闻字符串仅含标题新闻关键词分类code与名称

100 民生 故事 news_story
101 文化 文化 news_culture
102 娱乐 娱乐 news_entertainment
103 体育 体育 news_sports
104 财经 财经 news_finance
106 房产 房产 news_house
107 汽车 汽车 news_car
108 教育 教育 news_edu 
109 科技 科技 news_tech
110 军事 军事 news_military
112 旅游 旅游 news_travel
113 国际 国际 news_world
114 证券 股票 stock
115 农业 三农 news_agriculture
116 电竞 游戏 news_game

GitHub - aceimnorstuvwxz/toutiao-text-classfication-dataset: 今日头条中文新闻文本分类数据集

3自定义文本数据集

如果需要新增类别数据或者需要自定数据集进行训练可以如下进行处理

  • Train和Test数据集一个样本一个txt文本要求相同类别的文本放在同一个文件夹下且子目录文件夹命名为类别名称如

  • 类别文件 class_name.txt : (一行一个列表,最后一行,请多回车一行)
A
B
C
D
 
  • 修改配置文件数据路径config.yaml
# 训练数据集可支持多个数据集
train_data:
  - "data/dataset/train"
# 测试数据集
test_data:
  - "data/dataset/test"
vocab_file: "./data/dataset/vocabulary.json" # 字典文件(会根据训练数据集自动生成)
# 类别文件
class_name: "data/dataset/class_name.txt"

三、TextCNN模型结构

1TextCNN模型结构

TextCNN文本分类的网络结如下图所示可以分为4部分分别为输入层CNN层池化层和输出层

在这里插入图片描述

 以中文文本情感分类(二分类)作为简单的例子。

  1. 输入层也称embedding层TextCNN的输入序列是一个固定长度的句子图示中是由11个词组成一条句子context_size=11每个词用6维词向量表示(embedding_dim=6),即输入通道数in_channels=6。因此输入序列shape=(11,6)加上Batch这个维度则是shape=(batch_size,context_size,embedding_dim)=(B,11,6)
  2. CNN层也称卷积层由一维卷积核Conv1d组成左边的一维卷积核大小为2kernel_size=2输出通道数分别设为4右边的一维卷积核大小为4kernel_size=4输出通道数分别设为5卷积步长stride=1因此一维卷积计算后左边一维卷积输出宽度=11−2+1=10右边边一维卷积输出宽度11−4+1=8。
  3. 池化层将CNN层的输出的9个通道经过时序最大池化max_pool1d并将池化输出cat连结成一个9维向量。
  4. 分类层也是输出层由简单的全连接层组成对于简单二分类其输出维度2即正面情感和负面情感的预测概率。

2TextCNN实现

根据TextCNN网络结构我们可以使用Pytorch构建一个TextCNN模型

# -*-coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F


class GlobalMaxPool1d(nn.Module):
    def __init__(self):
        super(GlobalMaxPool1d, self).__init__()

    def forward(self, x):
        return F.max_pool1d(x, kernel_size=x.shape[2])  # shape: (batch_size, channel, 1)


class TextCNN(nn.Module):
    def __init__(self, num_classes, num_embeddings=-1, embedding_dim=128, kernel_sizes=[3, 4, 5, 6],
                 num_channels=[256, 256, 256, 256], embeddings_pretrained=None):
        """
        :param num_classes: 输出维度(类别数num_classes)
        :param num_embeddings: size of the dictionary of embeddings,词典的大小(vocab_size),
                               当num_embeddings<0,模型会去除embedding层
        :param embedding_dim:  the size of each embedding vector词向量特征长度
        :param kernel_sizes: CNN层卷积核大小
        :param num_channels: CNN层卷积核通道数
        :param embeddings_pretrained: embeddings pretrained参数默认None
        :return:
        """
        super(TextCNN, self).__init__()
        self.num_classes = num_classes
        self.num_embeddings = num_embeddings
        # embedding层
        if self.num_embeddings > 0:
            # embedding之后的shape: torch.Size([200, 8, 300])
            self.embedding = nn.Embedding(num_embeddings, embedding_dim)
            if embeddings_pretrained is not None:
                self.embedding = self.embedding.from_pretrained(embeddings_pretrained, freeze=False)
        # 卷积层
        self.cnn_layers = nn.ModuleList()  # 创建多个一维卷积层
        for c, k in zip(num_channels, kernel_sizes):
            cnn = nn.Sequential(
                nn.Conv1d(in_channels=embedding_dim,
                          out_channels=c,
                          kernel_size=k),
                nn.BatchNorm1d(c),
                nn.ReLU(inplace=True),
            )
            self.cnn_layers.append(cnn)
        # 最大池化层
        self.pool = GlobalMaxPool1d()
        # 输出层
        self.classify = nn.Sequential(
            nn.Dropout(p=0.2),
            nn.Linear(sum(num_channels), self.num_classes)
        )

    def forward(self, input):
        """
        :param input:  (batch_size, context_size, embedding_size(in_channels))
        :return:
        """
        if self.num_embeddings > 0:
            # 得到词嵌入(b,context_size)-->(b,context_size,embedding_dim)
            input = self.embedding(input)
            # (batch_size, context_size, channel)->(batch_size, channel, context_size)
        input = input.permute(0, 2, 1)
        y = []
        for layer in self.cnn_layers:
            x = layer(input)
            x = self.pool(x).squeeze(-1)
            y.append(x)
        y = torch.cat(y, dim=1)
        out = self.classify(y)
        return out


if __name__ == "__main__":
    device = "cuda:0"
    batch_size = 4
    num_classes = 2  # 输出类别
    context_size = 7  # 句子长度字词个数
    num_embeddings = 1024  # 词典的大小(vocab_size)
    embedding_dim = 6  # 词向量特征长度
    kernel_sizes = [2, 4]  # CNN层卷积核大小
    num_channels = [4, 5]  # CNN层卷积核通道数
    input = torch.ones(size=(batch_size, context_size)).long().to(device)
    model = TextCNN(num_classes=num_classes,
                    num_embeddings=num_embeddings,
                    embedding_dim=embedding_dim,
                    kernel_sizes=kernel_sizes,
                    num_channels=num_channels,
                    )
    model = model.to(device)
    model.eval()
    output = model(input)
    print("-----" * 10)
    print(model)
    print("-----" * 10)
    print(" input.shape:{}".format(input.shape))
    print("output.shape:{}".format(output.shape))
    print("-----" * 10)

测试模型打印结果


四、训练词嵌入word2vec可选

  • 不管是CNN还是RNN模型都是无法直接处理字符类别的单词因此需要通过某种方法把单词变成数字形式的向量才能作为模型的输入。把单词映射到向量空间中的一个向量的做法称为词嵌入word embedding对应的向量称为词向量word vector
  • 上面的TextCNN模型代码中定义了一个可学习的embedding层即词嵌入word2vec其作用就是将word序号ID转换为vector当然你也可以通过gensim训练自己的word2vec模型然后在数据处理中先将文本转换为词向量这样TextCNN就没有必要添加embedding层了。

项目仓库中提供了基于gensim的word2vec训练代码 word2vec.py 用户只需要修改好数据路径即可开始训练

# -*-coding: utf-8 -*-
"""
    @Author : panjq
    @E-mail : 390737991@qq.com
    @Date   : 2022-09-26 14:50:34
    @Brief  :
"""
import os
import sys

sys.path.insert(0, os.getcwd())
import random
import numpy as np
from gensim.models import word2vec
from core.utils import jieba_utils, nlp_utils
from pybaseutils import file_utils


class ChineseWord2Vector(object):
    """中文word2vec"""

    def __init__(self, stop_words=[], vector_size=128, window=5, min_count=5, epochs=10, workers=4):
        """
        :param stop_words: 停用词,用于ignore的字词
        :param vector_size: 是每个词的向量维度embedding_size
        :param window: 是词向量训练时的上下文扫描窗口大小窗口为5就是考虑前5个词和后5个词
        :param min_count: 设置最低频数默认是5如果一个词语在文档中出现的次数小于5那么就会丢弃
        :param epochs: Number of iterations (epochs) over the corpus. (Formerly: `iter`)
        :param workers:  是训练的线程数默认是当前运行机器的处理器核数
        """
        self.stop_words = stop_words if stop_words else jieba_utils.get_common_stop_words()
        self.vector_size = vector_size
        self.epochs = epochs
        self.window = window
        self.min_count = min_count
        self.workers = workers
        self.model: word2vec.Word2Vec = None

    def init_model(self):
        self.index_to_key = self.model.wv.index_to_key
        self.key_to_index = self.model.wv.key_to_index
        self.embedding = self.model.wv.vectors
        self.vector_size = self.model.wv.vector_size
        return self.model

    def cut_words_files(self, corpus: str, cutwords: str, user_file: str = "data/user_dict.txt", stop_words=[]):
        """
        :param corpus: 语料文件
        :param cutwords: jieba分词后保存的根目录
        :param user_file: 用户自定义的文件
        :param stop_words: 停用词,用于ignore的字词
        :return:
        """
        jieba_utils.load_userdict(user_file)
        print("corpus root     :{}".format(corpus))
        print("output cutwords :{}".format(cutwords))
        print("user_file       :{}".format(user_file))
        print("stop_words      :{}".format(stop_words))
        if not stop_words: stop_words = self.stop_words
        self.stop_words = stop_words
        nlp_utils.get_files_sentences_cutword(corpus, cutwords, stop_words=stop_words, block_size=10000)
        # 若只有一个文件使用LineSentence读取文件
        # sentences = word2vec.LineSentence(segment_path)
        # 若存在多文件使用PathLineSentences读取文件列表
        # sentences = word2vec.PathLineSentences(cutwords)
        sentences = word2vec.PathLineSentences(cutwords)
        return sentences

    def start_train(self, sentences):
        """
        :param sentences: *.txt文件路径所有字词需要预处理并被空格分隔
                        sentences可以是LineSentence或者PathLineSentences读取的文件对象也可以是
                        The `sentences` iterable can be simply a list of lists of tokens,
                        如lists=[['我','是','中国','人'],['我','的','家乡','在','广东']]

        """
        self.model = word2vec.Word2Vec(sentences,
                                       vector_size=self.vector_size,
                                       window=self.window,
                                       min_count=self.min_count,
                                       workers=self.workers,
                                       epochs=self.epochs,
                                       seed=2020,
                                       )

    def save_model(self, model_file) -> word2vec.Word2Vec:
        file_utils.create_file_path(model_file)
        self.model.save(model_file)
        self.init_model()
        return self.model

    def load_model(self, model_file) -> word2vec.Word2Vec:
        self.model = word2vec.Word2Vec.load(model_file)
        self.init_model()
        return self.model

    def get_similarity(self, key1, key2):
        """Compute cosine similarity between two keys."""
        return self.model.wv.similarity(key1, key2)

    def get_index(self, key, default=None):
        """Return the integer index (slot/position) where the given key's vector is stored in the backing vectors array."""
        return self.model.wv.get_index(key, default=default)

    def get_vector(self, key, norm=False):
        """Get the key's vector, as a 1D numpy array."""
        return self.model.wv.get_vector(key, norm=norm)

    def get_text_vector(self, text, context_size=-1, pad_token='<pad>'):
        """
        将句子中的所有词转为词向量
        :param text:
        :return: context_size 句子最大长度max_size
        :return: pad_token 句子不足时是否填充0
        """
        if context_size > 0: text = text[0:min(6 * context_size, len(text))]
        words = jieba_utils.cut_content_word(text, stop_words=self.stop_words)
        words = jieba_utils.padding_words(words, context_size=context_size, pad_token=pad_token)
        vector = self.get_words_vector(words)
        return vector

    def get_words_vector(self, words):
        """
        将word转换为vecror
        :param words:
        :return:
        """
        vector = []
        for w in words:
            try:
                v = self.get_vector(w)
            except Exception as e:
                v = np.zeros(shape=(self.model.vector_size,), dtype=np.float32)
            vector.append(v)
        vector = np.asarray(vector, dtype=np.float32)
        return vector

    def get_words_vector_padding(self, words, context_size=256, random_crop=False, padding=True):
        vector = []
        for w in words:
            try:
                v = self.get_vector(w)
                vector.append(v)
            except Exception as e:
                pass
        if len(vector) == 0: return []
        vector = np.asarray(vector, dtype=np.float32)
        nums, dims = vector.shape
        pad = context_size - nums
        if padding and pad > 0:
            zeros = np.zeros(shape=(pad, dims), dtype=np.float32)
            vector = np.concatenate([vector, zeros], axis=0)

        if random_crop and pad < 0:
            start = random.randint(0, nums - context_size)
            vector = vector[start:start + context_size, :]
        else:
            vector = vector[0:context_size, :]
        return vector


def train_simple_demo():
    source = './data/source'  # 文本数据路径
    user_file = 'data/user_dict.txt'
    cutwords = os.path.join(os.path.dirname(source), "cutwords")  # 分词结果
    model_file = os.path.join(os.path.dirname(source), "word2vec", "simple_word2vec128.model")
    wv_trainer = ChineseWord2Vector(vector_size=128, window=10, min_count=5, epochs=10)
    sentences = wv_trainer.cut_words_files(source, cutwords, user_file=user_file)
    wv_trainer.start_train(sentences)
    wv_trainer.save_model(model_file)
    model = wv_trainer.load_model(model_file)
    print("save word2vec:{}".format(model_file))
    # 测试
    w1 = '沙瑞金'
    w2 = '高育良'
    w3 = '车'
    vector = wv_trainer.get_vector(w1)
    print("({},{}),similarity={}".format(w1, w2, model.wv.similarity(w1, w2)))
    print("({},{}),similarity={}".format(w1, w3, model.wv.similarity(w1, w3)))
    # print("{} shape={},vector= \n{}".format(w1, vector.shape, vector))
    vector = wv_trainer.get_text_vector("我是一名中国人zhongguo")


def train_THUCNews():
    source = '/home/dm/nasdata/dataset/csdn/Text/THUCNews'  # 文本数据路径
    user_file = "./data/user_dict.txt"
    cutwords = os.path.join(os.path.dirname(source), "THUCNews-cutwords")  # 分词结果
    model_file = os.path.join(os.path.dirname(source), "word2vec128.model")
    wv_trainer = ChineseWord2Vector(vector_size=128, window=10, min_count=5, epochs=10)
    sentences = wv_trainer.cut_words_files(source, cutwords, user_file=user_file)
    wv_trainer.start_train(sentences)
    wv_trainer.save_model(model_file)
    model = wv_trainer.load_model(model_file)
    print("save word2vec:{}".format(model_file))
    # 测试
    w1 = '北京'
    w2 = '上海'
    w3 = '吃饭'
    vector = wv_trainer.get_vector(w1)
    print("({},{}),similarity={}".format(w1, w2, model.wv.similarity(w1, w2)))
    print("({},{}),similarity={}".format(w1, w3, model.wv.similarity(w1, w3)))
    # print("{} shape={},vector= \n{}".format(w1, vector.shape, vector))
    vector = wv_trainer.get_text_vector("我是一名中国人zhongguo")


if __name__ == '__main__':
    # 简单的训练词嵌入模型
    train_simple_demo()
    # 使用THUCNews数据训练词嵌入模型
    # train_THUCNews()

样例中使用小说《人民名义》 训练一个word2vec模型训练完成后测试单词(沙瑞金,高育良)的相似性similarity=0.8832而(沙瑞金,车)的相似性similarity=0.4969。


五、文本预处理

Pytorch的提供文本处理工具torchtext该工具功能非常强大提供了很多nlp方面的数据集可以直接加载使用也提供了不少训练好的词向量之类的但该工具封装的太高级了实际使用起来限制也太多了灵活性不高导致这个模块使用起来特别的别扭。所有后面干脆自己写Dataset数据处理方式了

对于中文文本数据预处理主要有两部分句子分词处理(英文文本不需要分词)特殊字符处理

1句子分词处理jieba中文分词

本博客使用jieba工具进行中文分词工具比较简单就不单独说明了安装方法

pip install jieba 

2特殊字符处理

jieba分词后会出现很多特殊字符需要进一步做一些的处理

  • 一些换行符空格等特殊字符以及一些标点符号。《》等这些特殊的字符称为stop_words需要剔除
  • 一些英文字母大小需要转换统一为小写
  • 一些繁体字统一转换为简体字等
  • 一些专有名词比如地名人名这些分词时需要整体切词jieba.load_userdict(file)

3文本数据增强

在计算机视觉图像识别任务中图像数据增强主要有裁剪、翻转、旋转、⾊彩变换等⽅式其目的增加数据的多样性提高模型的泛化能力。但是NLP任务中的数据是离散的无法像操作图片一样连续的方式操作文字这导致我们⽆法对输⼊数据进⾏直接简单地转换换掉⼀个词就有可能改变整个句⼦的含义。

常用的NLP文本数据增强方法主要有

  • 随机截取 随机截取文本一个片段
  • 同义词替换SR: Synonyms Replace不考虑stopwords在句⼦中随机抽取n个词然后从同义词词典中随机抽取同义词并进⾏替换。
  • 随机插⼊(RI: Randomly Insert)不考虑stopwords随机抽取⼀个词然后在该词的同义词集合中随机选择⼀个插⼊原句⼦中的随机位置。
  • 随机交换(RS: Randomly Swap)句⼦中随机选择两个词位置交换。
  • 随机删除(RD: Randomly Delete)句⼦中的每个词以概率p随机删除

项目已经实现随机截取随机插⼊随机删除等几种文本数据增强方式

# -*- coding: utf-8 -*-

import math
import random
from typing import List


def random_text_crop(text: List, label, context_size, token="<pad>", p=0.5):
    """
    句⼦中的每个词以概率p随机截取
    :param text:
    :param label:
    :param context_size:
    :param token:
    :param p:
    :return:
    """
    context_size = int(context_size)
    nums = len(text)
    pad = context_size - nums
    if pad > 0 and token:
        text = [token] * pad + text
    if random.random() < p and pad < 0:
        start = random.randint(0, nums - context_size)
        text = text[start:start + context_size]
    elif len(text) > context_size:
        text = text[0:context_size]
    return text, label


def random_text_mask(text: List, label, len_range=(0, 2), token="<pad>", p=0.5):
    """
    句⼦中的每个词以概率p替换成token
    :param text:
    :param label:
    :param len_range:
    :param p:
    :return:
    """
    if random.random() < p and len(text) > 2 * len_range[1]:
        nums = math.ceil(random.uniform(len_range[0], len_range[1]))
        for i in range(nums):
            index = int(random.uniform(0, len(text) - 1))
            text[index] = token
    return text, label


def random_text_delete(text: List, label, len_min, p=0.5):
    """
    句⼦中的每个词以概率p随机删除
    :param text:
    :param label:
    :param len_min: 句子最小长度低于该值不会删除
    :param p:
    :return:
    """
    if random.random() < p and len(text) > len_min:
        nums = int(random.uniform(0, len(text) - len_min))
        for i in range(nums):
            index = int(random.uniform(0, len(text)))
            del text[index]
    return text, label


def random_text_insert(text: List, label, len_range=(0, 2), token="<pad>", p=0.5):
    """
    句⼦中的每个词以概率p随机插入
    :param text:
    :param label:
    :param len_range: 
    :param p:
    :return:
    """
    if random.random() < p and len(text) > 2 * len_range[1]:
        nums = math.ceil(random.uniform(len_range[0], len_range[1]))
        for i in range(nums):
            index = int(random.uniform(0, len(text) - 1))
            text.insert(index, token)
    return text, label


if __name__ == '__main__':
    label = 1
    context_size = 10
    pad_token = "<pad>"
    p = 10
    for i in range(10):
        text = "我是一名中国人我爱中国我的家乡在广东"
        text = "_".join(text).split("_")
        len_range = (0, context_size // 4)
        # text, label = random_text_crop(text, label, 1.8 * context_size, token=None, p=0.8)
        # text, label = random_text_delete(text, label, len_min=1.5 * context_size)
        text, label = random_text_insert(text, label, len_range=len_range, token=pad_token)
        # text, label = random_text_mask(text, label, len_range=len_range, token=pad_token)
        # text, label = random_text_crop(text, label, context_size, token=pad_token, p=0.8)
        print(text, len(text))

六、训练过程 

项目以THUCNews文本分类数据集为作为训练数据训练一个基于TextCNN的文本分类模型这里为了简单没有使用gensim训练word2vec词向量模型而是在TextCNN模型代码中定义了一个可学习的embedding层用于代替word2vec

1项目框架说明

.
├── configs              # 训练配置文件
├── core                 # 模型和训练相关工具
├── data                 # 相关数据
├── modules              # 相关依赖包模块
├── work_space           # 训练模型输出文件目录
├── README.md            # 项目工程说明文档
├── requirements.txt     # 相关依赖包版本说明请用pip安装
├── word2vec.py          # 训练词嵌入模型
├── classifier.py        # 测试文本分类脚本
└── train.py             # 训练文件

项目依赖的python包请使用pip安装对应版本

numpy==1.16.3
matplotlib==3.1.0
Pillow==6.0.0
easydict==1.9
opencv-contrib-python==4.5.2.52
opencv-python==4.5.1.48
pandas==1.1.5
PyYAML==5.3.1
scikit-image==0.17.2
scikit-learn==0.24.0
scipy==1.5.4
seaborn==0.11.2
tensorboard==2.5.0
tensorboardX==2.1
torch==1.7.1+cu110
torchvision==0.8.2+cu110
tqdm==4.55.1
xmltodict==0.12.0
basetrainer
pybaseutils==0.6.9
jieba==0.42.1
gensim==4.2.0

2准备Train和Test文本数据

下载THUCNews文本数据集并解压由于原始数据没有划分训练集和测试集需要自己手动划分项目随机抽取每类的100张文本作为测试集其余的为训练集

然后根据自己的保存的数据路径修改配置文件数据路径config_textfolder.yaml

# 训练数据集可支持多个数据集
train_data:
  - "/path/to/dataset/THUCNews/train"
# 测试数据集
test_data:
  - "/path/to/dataset/THUCNews/test"
vocab_file: "./data/vocabulary/vocabulary.json" # 字典文件(会根据训练数据集自动生成),或者word2vec文件
# 类别文件
class_name: "path/to/dataset/THUCNews/class_name.txt"

3配置文件config_textfolder.yaml

# 训练数据集可支持多个数据集
train_data:
  - "/path/to/dataset/THUCNews/train"
# 测试数据集
test_data:
  - "/path/to/dataset/THUCNews/test"
vocab_file: "./data/vocabulary/vocabulary.json" # 字典文件(会根据训练数据集自动生成),或者word2vec文件
# 类别文件
class_name: "path/to/dataset/THUCNews/class_name.txt"


data_type: "textfolder"        # 加载数据DataLoader方法:word2vec,textfolder
flag: ""                       # 输出目录标识
resample: True                 # 是否进行重采样
work_dir: "work_space"         # 保存输出模型的目录
net_type: "TextCNN"            # 骨干网络,支持TextCNN,TextCNNv2,LSTM,BiLSTM等
context_size: 300              # 句子长度
topk: [ 1, ]                   # 计算topK的准确率
batch_size: 128                # 批训练大小
lr: 0.001                      # 初始学习率
optim_type: "Adam"             # 选择优化器SGD,Adam
loss_type: "CELoss"            # 选择损失函数支持CrossEntropyLoss(CELoss)
momentum: 0.9                  # SGD momentum
num_epochs: 160                # 训练循环次数
num_workers: 12                # 加载数据工作进程数
weight_decay: 0.00005          # weight_decay默认5e-4
#weight_decay: 0.0             # weight_decay默认5e-4
scheduler: "multi-step"        # 学习率调整策略
milestones: [ 90,120,140 ]     # 下调学习率方式
gpu_id: [ 0,1 ]                # GPU ID
log_freq: 10                   # LOG打印频率
pretrained: True               # 是否使用pretrained模型
finetune: False                # 是否进行finetune
  • 目标支持模型主要有TextCNN,LSTM,BiLSTM等详见模型等 其他模型可以自定义添加
  • 训练参数可以通过config.yaml配置文件
参数类型参考值说明
train_datastr, list-训练数据文件可支持多个文件
test_datastr, list-测试数据文件可支持多个文件
vocab_filestr-
字典文件(会根据训练数据集自动生成),或者word2vec文件
class_namestr-类别文件
data_typestr-加载数据DataLoader方法
resampleboolTrue是否进行重采样
work_dirstrwork_space训练输出工作空间
net_typestrTextCNN骨干网络,支持TextCNN,LSTM,BiLSTM等
context_sizeint128句子长度
topklist[1,3,5]计算topK的准确率
batch_sizeint32批训练大小
lrfloat0.1初始学习率大小
optim_typestrSGD优化器{SGD,Adam}
loss_typestrCELoss损失函数
schedulerstrmulti-step学习率调整策略{multi-step,cosine}
milestoneslist[30,80,100]降低学习率的节点仅仅scheduler=multi-step有效
momentumfloat0.9SGD动量因子
num_epochsint120循环训练的次数
num_workersint12DataLoader开启线程数
weight_decayfloat5e-4权重衰减系数
gpu_idlist[ 0 ]指定训练的GPU卡号可指定多个
log_freqint20显示LOG信息的频率
finetunestrmodel.pthfinetune的模型

4开始训练

整套训练代码非常简单操作用户只需要将相同类别的数据放在同一个目录下并填写好对应的数据路径即可开始训练了。

  • 如果你想验证项目可不可以训练请运行下面命令开始训练项目自带了小批量的文本数据方便测试项目代码正确情况下可以获得99%的文本分类准确率
python train.py -c configs/config.yaml 
  • 如果你想正式在THUCNews数据集上训练TextCNN文本分类模型请运行
python train.py -c configs/config_textfolder.yaml

以下是训练代码

# -*-coding: utf-8 -*-
"""
    @Author : panjq
    @E-mail : 390737991@qq.com
    @Date   : 2022-09-26 14:50:34
    @Brief  :
"""
import os
import torch
import argparse
import torch.nn as nn
import numpy as np
import tensorboardX as tensorboard
from tqdm import tqdm
from torch.utils import data as data_utils
from core.dataloader import build_dataset
from core.models import build_models
from core.criterion.build_criterion import get_criterion
from core.utils import torch_tools, metrics, log
from pybaseutils import file_utils, config_utils
from pybaseutils.metrics import class_report


class Trainer(object):
    def __init__(self, cfg):
        torch_tools.set_env_random_seed()
        # 设置输出路径
        time = file_utils.get_time()
        flag = [n for n in [cfg.net_type, cfg.loss_type, cfg.flag, time] if n]
        cfg.work_dir = os.path.join(cfg.work_dir, "_".join(flag))
        cfg.model_root = os.path.join(cfg.work_dir, "model")
        cfg.log_root = os.path.join(cfg.work_dir, "log")
        file_utils.create_dir(cfg.work_dir)
        file_utils.create_dir(cfg.model_root)
        file_utils.create_dir(cfg.log_root)
        file_utils.copy_file_to_dir(cfg.config_file, cfg.work_dir)
        config_utils.save_config(cfg, os.path.join(cfg.work_dir, "setup_config.yaml"))
        self.cfg = cfg
        self.topk = self.cfg.topk
        # 配置GPU/CPU运行设备
        self.gpu_id = cfg.gpu_id
        self.device = torch.device("cuda:{}".format(cfg.gpu_id[0]) if torch.cuda.is_available() else "cpu")
        # 设置Log打印信息
        self.logger = log.set_logger(level="debug", logfile=os.path.join(cfg.log_root, "train.log"))
        # 构建训练数据和测试数据
        self.train_loader = self.build_train_loader()
        self.test_loader = self.build_test_loader()
        # 构建模型
        self.model = self.build_model()
        # 构建损失函数
        self.criterion = self.build_criterion()
        # 构建优化器
        self.optimizer = self.build_optimizer()
        # 构建学习率调整策略
        self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, cfg.milestones)
        # 使用tensorboard记录和可视化Loss
        self.writer = tensorboard.SummaryWriter(cfg.log_root)
        # 打印信息
        self.num_samples = len(self.train_loader.sampler)
        self.logger.info("=" * 60)
        self.logger.info("work_dir          :{}".format(cfg.work_dir))
        self.logger.info("config_file       :{}".format(cfg.config_file))
        self.logger.info("gpu_id            :{}".format(cfg.gpu_id))
        self.logger.info("main device       :{}".format(self.device))
        self.logger.info("num_samples(train):{}".format(self.num_samples))
        self.logger.info("num_classes       :{}".format(cfg.num_classes))
        self.logger.info("mean_num          :{}".format(self.num_samples / cfg.num_classes))
        self.logger.info("=" * 60)

    def build_optimizer(self, ):
        """build_optimizer"""
        if self.cfg.optim_type.lower() == "SGD".lower():
            optimizer = torch.optim.SGD(params=self.model.parameters(), lr=self.cfg.lr,
                                        momentum=self.cfg.momentum, weight_decay=self.cfg.weight_decay)
        elif self.cfg.optim_type.lower() == "Adam".lower():
            optimizer = torch.optim.Adam(self.model.parameters(), lr=self.cfg.lr, weight_decay=self.cfg.weight_decay)
        else:
            optimizer = None
        return optimizer

    def build_train_loader(self, ) -> data_utils.DataLoader:
        """build_train_loader"""
        self.logger.info("build_train_loader,context_size:{}".format(self.cfg.context_size))
        dataset = build_dataset.load_dataset(data_type=self.cfg.data_type,
                                             filename=self.cfg.train_data,
                                             vocab_file=self.cfg.vocab_file,
                                             context_size=self.cfg.context_size,
                                             class_name=self.cfg.class_name,
                                             resample=self.cfg.resample,
                                             phase="train",
                                             shuffle=True)
        shuffle = True
        sampler = None
        self.logger.info("use resample:{}".format(self.cfg.resample))
        # if self.cfg.resample:
        #     weights = torch.DoubleTensor(dataset.classes_weights)
        #     sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, len(weights))
        #     shuffle = False
        loader = data_utils.DataLoader(dataset=dataset, batch_size=self.cfg.batch_size, sampler=sampler,
                                       shuffle=shuffle, num_workers=self.cfg.num_workers)
        self.cfg.num_classes = dataset.num_classes
        self.cfg.num_embeddings = dataset.num_embeddings
        self.cfg.class_name = dataset.class_name
        file_utils.copy_file_to_dir(self.cfg.vocab_file, cfg.work_dir)
        return loader

    def build_test_loader(self, ) -> data_utils.DataLoader:
        """build_test_loader"""
        self.logger.info("build_test_loader,context_size:{}".format(cfg.context_size))
        dataset = build_dataset.load_dataset(data_type=self.cfg.data_type,
                                             filename=self.cfg.test_data,
                                             vocab_file=self.cfg.vocab_file,
                                             context_size=self.cfg.context_size,
                                             class_name=self.cfg.class_name,
                                             phase="test",
                                             resample=False,
                                             shuffle=False)
        loader = data_utils.DataLoader(dataset=dataset, batch_size=self.cfg.batch_size,
                                       shuffle=False, num_workers=self.cfg.num_workers)
        self.cfg.num_classes = dataset.num_classes
        self.cfg.num_embeddings = dataset.num_embeddings
        self.cfg.class_name = dataset.class_name
        return loader

    def build_model(self, ) -> nn.Module:
        """build_model"""
        self.logger.info("build_model,net_type:{}".format(self.cfg.net_type))
        model = build_models.get_models(net_type=self.cfg.net_type,
                                        num_classes=self.cfg.num_classes,
                                        num_embeddings=self.cfg.num_embeddings,
                                        embedding_dim=128,
                                        is_train=True,
                                        )
        if self.cfg.finetune:
            self.logger.info("finetune:{}".format(self.cfg.finetune))
            state_dict = torch_tools.load_state_dict(self.cfg.finetune)
            model.load_state_dict(state_dict)
        model = model.to(self.device)
        model = nn.DataParallel(model, device_ids=self.gpu_id, output_device=self.device)
        return model

    def build_criterion(self, ):
        """build_criterion"""
        self.logger.info(
            "build_criterion,loss_type:{}, num_embeddings:{}".format(self.cfg.loss_type, self.cfg.num_embeddings))
        criterion = get_criterion(self.cfg.loss_type, self.cfg.num_embeddings, device=self.device)
        # criterion = torch.nn.CrossEntropyLoss()
        return criterion

    def train(self, epoch):
        """训练"""
        train_losses = metrics.AverageMeter()
        train_accuracy = {k: metrics.AverageMeter() for k in self.topk}
        self.model.train()  # set to training mode
        log_step = max(len(self.train_loader) // cfg.log_freq, 1)
        for step, data in enumerate(tqdm(self.train_loader)):
            inputs, target = data
            inputs, target = inputs.to(self.device), target.to(self.device)
            outputs = self.model(inputs)
            loss = self.criterion(outputs, target)
            self.optimizer.zero_grad()  # 反馈
            loss.backward()
            self.optimizer.step()  # 更新
            train_losses.update(loss.cpu().data.item())
            # 计算准确率
            target = target.cpu()
            outputs = outputs.cpu()
            outputs = torch.nn.functional.softmax(outputs, dim=1)
            pred_score, pred_index = torch.max(outputs, dim=1)
            acc = metrics.accuracy(outputs.data, target, topk=self.topk)
            for i in range(len(self.topk)):
                train_accuracy[self.topk[i]].update(acc[i].data.item(), target.size(0))
            if step % log_step == 0:
                lr = self.scheduler.get_last_lr()[0]  # 获得当前学习率
                topk_acc = {"top{}".format(k): v.avg for k, v in train_accuracy.items()}
                self.logger.info(
                    "train {}/epoch:{:0=3d},lr:{:3.4f},loss:{:3.4f},acc:{}".format(step, epoch, lr, train_losses.avg,
                                                                                    topk_acc))

        topk_acc = {"top{}".format(k): v.avg for k, v in train_accuracy.items()}
        self.writer.add_scalar("train-loss", train_losses.avg, epoch)
        self.writer.add_scalars("train-accuracy", topk_acc, epoch)
        self.logger.info("train epoch:{:0=3d},loss:{:3.4f},acc:{}".format(epoch, train_losses.avg, topk_acc))
        return topk_acc["top{}".format(self.topk[0])]

    def test(self, epoch):
        """测试"""
        test_losses = metrics.AverageMeter()
        test_accuracy = {k: metrics.AverageMeter() for k in self.topk}
        true_labels = np.ones(0)
        pred_labels = np.ones(0)
        self.model.eval()  # set to evaluates mode
        with torch.no_grad():
            for step, data in enumerate(tqdm(self.test_loader)):
                inputs, target = data
                inputs, target = inputs.to(self.device), target.to(self.device)
                outputs = self.model(inputs)
                loss = self.criterion(outputs, target)
                test_losses.update(loss.cpu().data.item())
                # 计算准确率
                target = target.cpu()
                outputs = outputs.cpu()
                outputs = torch.nn.functional.softmax(outputs, dim=1)
                pred_score, pred_index = torch.max(outputs, dim=1)
                acc = metrics.accuracy(outputs.data, target, topk=self.topk)
                true_labels = np.hstack([true_labels, target.numpy()])
                pred_labels = np.hstack([pred_labels, pred_index.numpy()])

                for i in range(len(self.topk)):
                    test_accuracy[self.topk[i]].update(acc[i].data.item(), target.size(0))

        report = class_report.get_classification_report(true_labels, pred_labels, target_names=self.cfg.class_name)
        topk_acc = {"top{}".format(k): v.avg for k, v in test_accuracy.items()}
        lr = self.scheduler.get_last_lr()[0]  # 获得当前学习率
        self.writer.add_scalar("test-loss", test_losses.avg, epoch)
        self.writer.add_scalars("test-accuracy", topk_acc, epoch)
        self.logger.info("test  epoch:{:0=3d},lr:{:3.4f},loss:{:3.4f},acc:{}".format(epoch, lr, test_losses.avg, topk_acc))
        self.logger.info("{}".format(report))
        return topk_acc["top{}".format(self.topk[0])]

    def run(self):
        """开始运行"""
        self.max_acc = 0.0
        for epoch in range(self.cfg.num_epochs):
            train_acc = self.train(epoch)  # 训练模型
            test_acc = self.test(epoch)  # 测试模型
            self.scheduler.step()  # 更新学习率
            lr = self.scheduler.get_last_lr()[0]  # 获得当前学习率
            self.writer.add_scalar("lr", lr, epoch)
            self.save_model(self.cfg.model_root, test_acc, epoch)
            self.logger.info("epoch:{}, lr:{}, train acc:{:3.4f}, test acc:{:3.4f}".
                             format(epoch, lr, train_acc, test_acc))

    def save_model(self, model_root, value, epoch):
        """保存模型"""
        # 保存最优的模型
        if value >= self.max_acc:
            self.max_acc = value
            model_file = os.path.join(model_root, "best_model_{:0=3d}_{:.4f}.pth".format(epoch, value))
            file_utils.remove_prefix_files(model_root, "best_model_*")
            torch.save(self.model.module.state_dict(), model_file)
            self.logger.info("save best   model file:{}".format(model_file))
        # 保存最新的模型
        name = "model_{:0=3d}_{:.4f}.pth".format(epoch, value)
        model_file = os.path.join(model_root, "latest_{}".format(name))
        file_utils.remove_prefix_files(model_root, "latest_*")
        torch.save(self.model.module.state_dict(), model_file)
        self.logger.info("save latest model file:{}".format(model_file))
        self.logger.info("-------------------------" * 4)


def get_parser():
    # cfg_file = "configs/config_textfolder.yaml"
    cfg_file = "configs/config.yaml"
    parser = argparse.ArgumentParser(description="Training Pipeline")
    parser.add_argument("-c", "--config_file", help="configs file", default=cfg_file, type=str)
    cfg = config_utils.parser_config(parser.parse_args(), cfg_updata=True)
    return cfg


if __name__ == "__main__":
    cfg = get_parser()
    train = Trainer(cfg)
    train.run()

5可视化训练过程

训练过程可视化工具是使用Tensorboard使用方法
# 基本方法
tensorboard --logdir=path/to/log/
# 例如请修改自己的训练的模型路径
tensorboard --logdir=work_space/TextCNN_CELoss_20230106152138/log

可视化效果 

​​​​
​​​​
​​​​

6一些优化建议

训练完成后目前基于TextCNN的文本分类识别在THUCNews数据集上训练集的Accuracy 99%左右测试集的Accuracy在88.36%左右如果想进一步提高准确率可以尝试

  1. 数据整合部分分类之间本身模棱两可例如体育和娱乐教育和科技本身类别就有很多相似之处导致模型分类困难THUCNews数据量虽然庞大但不是十分干净有很多脏数据建议你训练前清洗或整合部分数据集不然会影响模型的识别的准确率。
  2. 增加TextCNN参数量比如将TextCNN的num_channels设置大一点当然模型越复杂越容易过拟合
  3. 增加pretrained模型项目构建TextCNN模型随机初始化了一个可学习的二维矩阵Embedding该Embedding模型没有增加pretrained的若能加入pretrained其准确率会好很多。
  4. 文本数据增强如同义词替换文本随机插入随机删除等处理增强模型泛化能力
  5. 样本均衡数据不均衡部分类目数据太少 建议进行样本均衡处理减少长尾问题的影响
  6. 超参调优 比如学习率调整策略优化器SGD,Adam等
  7. 损失函数 目前训练代码已经支持交叉熵LabelSmoothing可以尝试FocalLoss等损失函数

七. 模型测试效果

classifier.py文件用于模型推理和测试脚本填写好配置文件模型文件以及测试文本路径即可运行测试了

def get_parser():
    model_file = "work_space/TextCNN_CELoss_20221226114529/model/latest_model_159_0.8714.pth"
    config_file = os.path.join(os.path.dirname(os.path.dirname(model_file)), "config_textfolder.yaml")
    vocab_file = os.path.join(os.path.dirname(os.path.dirname(model_file)), "vocabulary.json")
    text_dir = "data/test-text"
    parser = argparse.ArgumentParser(description="Inference Argument")
    parser.add_argument("-c", "--config_file", help="configs file", default=config_file, type=str)
    parser.add_argument("-m", "--model_file", help="model_file", default=model_file, type=str)
    parser.add_argument("-v", "--vocab_file", help="vocab_file", default=vocab_file, type=str)
    parser.add_argument("--device", help="cuda device id", default="cuda:0", type=str)
    parser.add_argument("--text_dir", help="text", default=text_dir, type=str)
    return parser

在项目根目录终端运行命令\表示换行符

#!/usr/bin/env bash
# Usage:
# python classifier.py -c "path/to/config.yaml" -m "path/to/model.pth" -v "path/to/vocabulary.json" --text_dir "path/to/text_dir"

python classifier.py \
  -c "work_space/TextCNN_CELoss_20221226114529/config_textfolder.yaml" \
  -m "work_space/TextCNN_CELoss_20221226114529/model/latest_model_159_0.8714.pth" \
  -v "work_space/TextCNN_CELoss_20221226114529/vocabulary.json" \
  --text_dir "data/test-text"

运行测试结果 


八.项目源码下载

整套项目源码下载Pytorch TextCNN实现中文文本分类(附完整训练代码)

整套项目源码内容包含

  • 提供中文文本数据集THUCNews
  • 项目支持训练词嵌入模型训练word2vec.py
  • 项目提供Pytorch版本的中文文本分类模型训练train.py支持TextCNN, LSTM, BiLSTM等模型
  • 提供中文文本分类预测classifier.py
  • 简单配置一键开启训练自己的中文文本分类模型

如果你想学习中文单词预测请参考《Pytorch LSTM实现中文单词预测(附完整训练代码)_

阿里云国内75折 回扣 微信号:monov8
阿里云国际,腾讯云国际,低至75折。AWS 93折 免费开户实名账号 代冲值 优惠多多 微信号:monov8 飞机:@monov6