Pytorch LSTM实现中文单词预测（附完整训练代码）

Pytorch LSTM实现中文单词预测（附完整训练代码）

目录

​​Pytorch LSTM实现中文单词预测(词语预测 附完整训练代码)​​

​​1、项目介绍​​

​​2、中文单词预测方法（N-Gram 模型）​​

​​3、训练词嵌入word2vec（可选）​​

​​4、文本预处理​​

​​（1）句子分词处理：jieba中文分词​​

​​（2）特殊字符处理​​

​​（3）文本数据增强​​

​​（4）样本均衡（重点）​​

​​5、训练过程 ​​

​​（1）项目框架说明​​

​​（2）准备文本数据​​

​​（3）配置文件：config.yaml​​

​​（4）开始训练​​

​​（6）一些优化建议​​

​​6. 模型测试效果​​

​​7.项目源码下载​​

1、项目介绍

本文将分享一个NLP项目实例，实现一个类似于中文输入法中联想的功能；项目利用深度学习框架Pytorch，构建一个LSTM（也支持NGram，TextCNN，LSTM，BiLSTM等）模型，实现一个简易的中文单词预测（词语预测）功能，该功能可以根据用户输入的中文语句，自动预测(补充)词语；基于该项目训练的中文单词预测(词语预测）模型，在自定义的数据集上Top-1准确率最高可以达到91%左右，Top-5准确率最高可以达到97%左右。

如果，你想学习NLP中文文本分类，可参考另一篇博文《​​Pytorch TextCNN实现中文文本分类(附完整训练代码)​​》

2、中文单词预测方法（N-Gram 模型）

首先简单介绍一下 N-Gram 模型的原理：对于一英文句话，单词的排列顺序是非常重要的，所以我们能否由前面的几个词来预测后面的几个单词呢，比如 'I lived in France for 10 years, I can speak _' 这句话中，我们能够预测出最后一个词是 French。

对于一句话T，其由w1,w2....wn这n个词构成，可以得到下面的公式：

但是这样的一个模型参数过大，预测一个词需要前面所有的词作为条件来计算概率。我们可以再简化一下这个模型，比如对于一个词，并不需要前面所有的词作为条件概率，也就是说一个词可以只与其前面的几个词有关，这就是马尔科夫假设。

对于这里的条件概率，传统的方法是统计语料中每个词出现的频率，根据贝叶斯定理来估计这个条件概率，这里我们就可以用词嵌入对其进行代替，然后使用 RNN 进行条件概率的计算，然后最大化这个条件概率不仅修改词嵌入，同时能够使得模型可以依据计算的条件概率对其中的一个单词进行预测。

类似的，对于中文，我们也可以这样进行处理，比如一个简单的语句：【我是一名中国人】，我们希望如果我们输入【我是一名】，模型输出结果是【中国人】；只不过中文不像英文那样有明显空格作为单词分隔符，中文语句需要我们自己按照一定规则进行字词分割，这个工具可以使用jieba中文分词工具。

下面定义一个简单的NGram模型：

# -*-coding: utf-8 -*-
"""
    @Author : 390737991
    @E-mail : [email protected]
    @Date   : 2022-11-01 17:54:33
    @Brief  :
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class NGram(nn.Module):
    """N-Gram模型m"""

    def __init__(self, num_classes, context_size, num_embeddings, embedding_dim=128, embeddings_pretrained=None):
        """
        :param num_classes: 输出维度(类别数num_classes)
        :param context_size: 句子长度
        :param num_embeddings: size of the dictionary of embeddings,词典的大小(vocab_size)
        :param embedding_dim:  the size of each embedding vector，词向量特征长度
        :param embeddings_pretrained: embeddings pretrained参数，默认None
        :return:
        """
        self.num_classes = num_classes
        self.num_embeddings = num_embeddings
        super(NGram, self).__init__()
        # 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.classify = nn.Sequential(
            nn.Dropout(p=0.2),
            nn.Linear(embedding_dim * context_size, 512),
            nn.ReLU(True),
            nn.Dropout(p=0.2),
            nn.Linear(512, self.num_classes)
        )

    def forward(self, x):
        if self.num_embeddings > 0:
            x = self.embedding(x)  # 得到词嵌入
        x = x.view(x.size(0), -1)  # 将两个词向量拼在一起
        out = self.classify(x)
        return out


if __name__ == "__main__":
    batch_size = 2
    num_classes = 100
    num_embeddings = num_classes  # 预测的类别数目和单词数目一样
    context_size = 8  # 句子长度，即最大依赖的单词数目
    input = torch.ones(batch_size, context_size).long().cuda()
    model = NGram(num_classes, context_size, num_embeddings=num_embeddings, embedding_dim=64).cuda()
    print(model)
    out = model(input)
    print(out)
    print("input", input.shape)
    print("out  ", out.shape)

上面我们定义了一个NGram模型，其中参数 context_size 表示句子最大长度，表示我们希望由前面几个单词来预测这个单词，这里context_size=8个，表示由8单词（不足8个可以填充）预测1个单词；embedding_dim 表示词嵌入的维度，即词向量特征长度。num_embeddings是词典的大小，即我们字库的个数大小，由于输出预测的类别数等于词典的字词数的大小，所num_embeddings = num_classes；

NGram模型比较简单，如果去除Embedding层，其实就是一个由2层全连接层构成的模型；在实际应用中，效果比较差的，只能作为介绍使用；后续项目，将以LSTM模型为例，进行训练和测试，LSTM模型定义如下，其中我增加LayerNorm用于归一化数据，其作用类似于CNN中的BatchNorm。

import torch
import torch.nn as nn
import torch.nn.functional as F


class LayerNorm(nn.Module):
    """nn.LayerNorm"""

    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))

    def forward(self, input):
        mu = torch.mean(input, dim=-1, keepdim=True)
        sigma = torch.std(input, dim=-1, keepdim=True).clamp(min=self.eps)
        output = (input - mu) / sigma
        return output * self.weight.expand_as(output) + self.bias.expand_as(output)


class LSTMNet(nn.Module):
    def __init__(self, num_classes, num_embeddings, embedding_dim=64, hidden_size=128, num_layers=2,
                 batch_first=True, bidirectional=False, dropout=0, use_norm=True, embeddings_pretrained=None):
        """
        https://github.com/ne7ermore/torch-light/blob/master/lstm-text-classfication/model.py
        input为（batch,context_size,input_size）=（batch,context_size,dim_size）
        :param num_classes: 输出维度(类别数num_classes)
        :param num_embeddings: size of the dictionary of embeddings,词典的大小(vocab_size)
        :param embedding_dim: 输入特征的数目，输入向量维数
        :param hidden_size: 隐层的特征数目
        :param num_layers: 这个是模型集成的LSTM的个数 记住这里是模型中有多少个LSTM摞起来 一般默认就1个
        :param batch_first:
        :param bidirectional: True 代表使用双向LSTM
        """
        super(LSTMNet, self).__init__()
        self.use_norm = use_norm
        self.num_classes = num_classes
        self.num_embeddings = num_embeddings
        # embedding层
        if self.num_embeddings > 0:
            self.embedding = nn.Embedding(num_embeddings, embedding_dim)  # embedding之后的shape: torch.Size([200, 8, 300])
            if embeddings_pretrained is not None:
                self.embedding = self.embedding.from_pretrained(embeddings_pretrained, freeze=False)

        num_directions = 2 if bidirectional else 1
        self.lstm = nn.LSTM(input_size=embedding_dim, hidden_size=hidden_size,
                            num_layers=num_layers, batch_first=batch_first,
                            bidirectional=bidirectional, dropout=dropout)
        if self.use_norm:
            self.layer_norm = LayerNorm(hidden_size * num_directions)
            # self.layer_norm = nn.LayerNorm(hidden_size * num_directions)
        # 输出层
        self.classify = nn.Sequential(
            nn.Dropout(p=0.2),
            nn.Linear(hidden_size * num_directions, self.num_classes)
        )

    def forward(self, x):
        """
        :param x:  x's shape (batch_size, 序列长度, 序列中每个数据的长度)
        :return:
        """
        if self.num_embeddings > 0:
            x = self.embedding(x)  # 得到词嵌入(b,context_size)-->(b,context_size,embedding_dim)
        x, _, = self.lstm(x)  # x's shape (batch_size, 序列长度, hidden_dim)
        if self.use_norm:
            x = self.layer_norm(x)
        x = x[:, -1, :]  # 中间的序列长度取-1，表示取序列中的最后一个数据，这个数据长度为hidden_dim，
        # 得到的out的shape为(batch_size, hidden_dim)
        x = self.classify(x)  # 经过线性层后，out的shape为(batch_size, n_class)
        return x

    def forward_bk(self, inputs):
        # x.shape : batch,context_size,hidden_size , hn.shape and cn.shape : num_layes * direction_numbers,batch,hidden_size
        out, (x, cell) = self.lstm(inputs)
        if self.use_norm:
            x = self.layer_norm(x)
        a, b, c = x.shape
        x = self.classify(x.reshape(a * b, c))
        return x


if __name__ == "__main__":
    import numpy

    # context_size, kernel_sizes, num_channels = 100, [3, 4, 5], [128, 64, 32]
    batch_size = 2
    num_embeddings = 128
    context_size = 8
    # input = torch.ones(batch_size, embedding_size, context_size).float().cuda()
    input = torch.ones(batch_size, context_size).long().cuda()
    model = LSTMNet(num_embeddings, embedding_dim=64).cuda()
    print(model)
    out = model(input)
    print(out)
    print("input", input.shape)
    print("out  ", out.shape)

3、训练词嵌入word2vec（可选）

不管是CNN还是RNN模型，都是无法直接处理字符类别的单词，因此我们需要对单词进行编码，即通过某种方法把单词变成数字形式的向量才能作为模型的输入。把单词映射到向量空间中的一个向量的做法称为词嵌入（word embedding），对应的向量称为词向量（word vector）

上面的NGram模型代码中，定义了一个可学习的embedding层，即词嵌入word2vec，其作用就是将word序号ID转换为vector；当然你也可以通过gensim训练自己的word2vec模型，然后在数据处理中先将文本转换为词向量，这样NGram就没有必要添加embedding层了。

4、文本预处理

接下来，我们需要将句子按照context_size+1的长度进行逐个截取，前context_size个单词是模型输入数据，最后一个是预测结果，这样就构建了我们的训练集，核心代码如下：

def get_item_list(self, sentences, context_size, stride=1, padding="<pad>", shuffle=False):
        """
        构建数据集
        :param sentences: 语料数据[list],一句话一个列表
        :param context_size: 句子最大长度
        :param stride: 步长，默认1
        :param padding: 不足context_size，进行填充
        :return:
        """
        item_list = []
        for content in sentences:
            pad_size = context_size + 1 - len(content)
            if pad_size > 0:
                content = [padding] * pad_size + content
            for i in range(0, len(content) - context_size, stride):
                inputs = content[i:(i + context_size)]
                target = content[i + context_size]
                item_list.append((inputs, target))
        if shuffle:
            random.seed(100)
            random.shuffle(item_list)
        return item_list

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

（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随机删除

（4）样本均衡（重点）

有一些常用词汇，由于其出现的频率很高，导致模型预测的结果，会偏向于预测高频率出现的词汇；比如【的】字，在中文语句中，出现的频率特别大，导致模型预测的时候，输出结果经常被预测为【的】，显示这是不符合实际情况；一种行之有效的解决方法，是对数据进行均衡采用，即高频词汇应该降低采样次数，而低频词应该增加其采用次数。

项目已经实现样本均衡算法，config.yaml配置文件中，只需要设置resample: True即可开启样本均衡训练

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

# -*- 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))

5、训练过程 

（1）项目框架说明

.
├── configs              # 训练配置文件
├── core                 # 模型和训练相关工具
├── data                 # 相关数据
├── modules              # 相关依赖包模块
├── work_space           # 训练模型输出文件目录
├── README.md            # 项目工程说明文档
├── requirements.txt     # 相关依赖包版本说明，请用pip安装
├── predictor.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）准备文本数据

首先，我们需要收集中文文本数据集，由于我们是做单词预测算法，要求训练数据尽可能干净；考虑到我们的模型比较简单，无需像BERT那样海量数据。作为简单的Demo，项目从百度文库中收集了一些中文造句的常用句子，大概5千字的数据量吧

 然后根据自己的保存的数据路径，修改配置文件数据路径：configs/config.yaml （项目的文本数据放在data/text/data"中，可自行增加补充数据）；考虑到，单词预测是一种比较模糊预测的任务，因此，项目没有严格区分训练集和测试集，而是将训练数据和测试集都使用同一数据集。

# 训练数据集，可支持多个数据集
train_data:
  - "data/text/data"
# 测试数据集
test_data:
  - "data/text/data"
vocab_file: "data/text/vocabulary.json" # 指定字典的路径(会根据训练数据集自动生成)

（3）配置文件：config.yaml

# 训练数据集，可支持多个数据集
train_data:
  - "data/text/data"
# 测试数据集
test_data:
  - "data/text/data"
vocab_file: "data/text/vocabulary.json" # 指定字典的路径(会根据训练数据集自动生成)
class_name: ""
data_type: "textdata"          # 加载数据DataLoader方法
flag: ""                       # 输出目录标识
resample: True                 # 是否进行重采样
work_dir: "work_space"         # 保存输出模型的目录
net_type: "LSTM"            # 骨干网络,支持：NGram,TextCNN,LSTM,BiLSTM等
context_size: 8                # 句子长度
topk: [ 1,3,5 ]                # 计算topK的准确率
batch_size: 64                 #  批训练大小
lr: 0.01                       # 初始学习率
optim_type: "Adam"             # 选择优化器，SGD,Adam
loss_type: "CELoss"            # 选择损失函数：支持CrossEntropyLoss(CELoss)
momentum: 0.9                  # SGD momentum
num_epochs: 120                # 训练循环次数
num_workers: 0                 # 加载数据工作进程数
weight_decay: 0.0005           # weight_decay，默认5e-4
scheduler: "multi-step"        # 学习率调整策略
milestones: [ 30,70,100 ]      # 下调学习率方式
gpu_id: [ 0 ]                  # GPU ID
log_freq: 20                   # LOG打印频率
pretrained: True               # 是否使用pretrained模型
finetune: False                # 是否进行finetune

• 目标支持模型主要有：NGram,TextCNN,LSTM,BiLSTM等，详见模型等 ，其他模型可以自定义添加
• 训练参数可以通过config.yaml配置文件

（4）开始训练

整套训练代码非常简单操作，用户只需要将文本数据放在项目"data/text/data"目录下（也可以自定义数据路径），并填写好对应的数据路径，即可开始训练了。

• 如果你想验证项目可不可以训练，请运行下面命令开始训练；项目自带了小批量的文本数据，方便测试项目代码；对于简单的样本数据集，可以获得95%左右的预测准确率

python train.py -c configs/config.yaml

•  如果你想训练自己的数据，请准备好文本数据，并放在data/text/data中，文本只支持TXT格式，不支持PDF和word文档格式
• 配置文件configs/config.yaml的参数net_type，用于选择骨干网络,可以填写NGram, TextCNN, LSTM, BiLSTM等，后面模型以LSTM模型为准

以下是训练代码：

# -*-coding: utf-8 -*-
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.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/BiLSTM_CELoss_20230110175943/log

可视化效果 

（6）一些优化建议

训练完成后，可查看其Top-1,Top-3和Top-5准确率，其中NGram的Top-1准确率约0.8630，而TextCNN的准确率约0.9065，LSTM的准确率约0.9088，BiLSTM的准确率最高可以达到0.9100

1. 数据整合：建议对数据进行去燥，删除一些语句不通的文本；
2. 由于数据集比较小，有很多中文字词是不支持，建议增大数据集进行训练；
3. 增加LSTM参数量：比如增大LSTM的个数
4. 增加pretrained模型：项目构建LSTM模型，随机初始化了一个可学习的二维矩阵：Embedding，该Embedding模型没有增加pretrained的，若能加入pretrained，其准确率会好很多。
5. 文本数据增强：如同义词替换，文本随机插入，随机删除等处理，增强模型泛化能力
6. 样本均衡：数据不均衡，部分类目数据太少； 建议进行样本均衡处理，减少长尾问题的影响
7. 超参调优： 比如学习率调整策略，优化器（SGD,Adam等）
8. 损失函数： 目前训练代码已经支持：交叉熵，LabelSmoothing，可以尝试FocalLoss等损失函数

6. 模型测试效果

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

def get_parser():
    model_file = "work_space/LSTM_CELoss_20230110175804/model/best_model_119_0.9088.pth"
    config_file = os.path.join(os.path.dirname(os.path.dirname(model_file)), "config.yaml")
    vocab_file = os.path.join(os.path.dirname(os.path.dirname(model_file)), "vocabulary.json")
    input = "美丽豪华的 获得优异的"  # 才能获得优异的成绩
    # 我们家住的楼上有许多只壁虎，每天晚上你都能看到它们。小壁虎身子是土黄色的，一双圆溜溜的眼睛，尖尖的脑袋，拖着一条长长的尾巴
    input = "我们家住的楼上有许多只/每天晚上你都能看到 一双圆溜溜的 拖着一条长长的"  # 才能获得优异的成绩
    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("--input", help="text", default=input, type=str)
    return parser

或者在项目根目录终端运行命令（\表示换行符）：

#!/usr/bin/env bash
python predictor.py \
  -c "work_space/LSTM_CELoss_20230110175804/config.yaml" \
  -m "work_space/LSTM_CELoss_20230110175804/model/best_model_119_0.9088.pth" \
  -v "work_space/LSTM_CELoss_20230110175804/vocabulary.json" \

使用方法： 

1. 【使用说明】:输入任意文本，用[空格]或[/]表示需要预测的字词;输入[e]退出程序
2. 【输入例子】:美丽豪华的/获得优异的
3. 【结果说明】:输出括号内表示预测结果

运行测试结果：  

• 输入："我们家住的楼上有许多只/每天晚上你都能看到 一双圆溜溜的 拖着一条长长的"，返回的预测结果：

• 输入"美丽豪华的 获得优异的" ，返回的预测结果：

7.项目源码下载

整套项目源码内容包含

• 提供中文文本数据集，用于模型训练：数据主要从百度文库中收集了一些中文造句的常用句子，大概5千字的数据量
• 提供Pytorch版本的中文单词预测模型训练工具：train.py，支持NGram，TextCNN, LSTM, BiLSTM等模型训练和测试
• 提供中文单词预测测试脚本：predictor.py
• 项目已经实现样本均衡算法，config.yaml配置文件中，只需要设置resample: True即可开启样本均衡训练
• 简单配置，一键开启训练自己的中文单词预测模型

如果，你想学习NLP中文文本分类，可参考另一篇博文《​​Pytorch TextCNN实现中文文本分类(附完整训练代码)​​》