pytorch从头开始进行文本分类

前言

最近在做实体抽取的时候，一篇文章大约有几千字，按照300字长度进行切割后，会生成数量不等的句子，若是句子少还行，句子多的情况下，则会对造成巨大的计算负担，因为一篇文章中存在关键词的段落是比较少的，为了减轻计算负担，让实体抽取模型仅对有实体的段落进行预测是最佳的选择。首先我是思考了前后各2个段落的方式进行句子筛选，然而偏偏有文章实体是出现在文章中间的，因此不得不考虑对段落进行筛选，采用关键词匹配的方式进行筛选通常都会产生多余的句子，还是无法解决计算负担的问题，因此采用模型的方式进行是最佳的。
文本分类的模型选择有很多，可以采用BERT系列的模型，用专门对于文本分类进行微调后的BERT模型固然可以达到一个比较好的精度，然而却也是增加了计算负担。好在对于实体是否存在这个判断仅仅是粗判断，因此采用CNN也是可以完成任务的，这才有了这篇文章的出现。由于我之前都是用keras进行模型的训练与部署，最近转用torch训练模型，转onnx进行部署的方式，并且明年torch2可能出现，并大幅提升训练与部署速度，因此记录一下基于torch的文本分类模型

词典与字典

若是用词作为最小单位，则生成的是词典。若是用字作为最小单位，则生成的是字典。对于中文来说，字典的性能与词典类似，并且词向量空间小，是个比较好的选择，但是通常预训练好的向量如word2vec均为词向量，因此有预训练向量的情况下，词向量也是不错的选择。首先对训练文本生成字典，字典的主要作用是将字映射成数字

import json
import re
import glob
import torch
import numpy as np
np.random.seed(42)
from collections import Counter
from torch.utils.data import Dataset, DataLoader
from torch import nn

def get_vocab():
    counter = Counter()
    with open(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/labeled_data_for_mrc.jsonl', 'r', encoding='utf8') as f:
        for i, line in enumerate(f):
            if i == 1000:
                break
            line = json.loads(line)
            text = re.sub('\s+', '', line['text'])
            counter.update(text)

    updated_files = glob.glob(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/updated*')
    for file in updated_files:
        with open(file, 'r', encoding='utf8') as f:
            for i, line in enumerate(f):
                line = json.loads(line)
                text = re.sub('\s+', '', line['text'])
                counter.update(text)
    vocabs = list(zip(*counter.most_common()))[0]
    with open('vocab.txt', 'w', encoding='utf8') as f:
        f.write('[PAD]\n')
        f.write('[UNK]\n')
        f.write('\n'.join(vocabs))
get_vocab()

由于我的文本比较多，因此看起来比较复杂，其实最主要的内容就是将文本都更新到Counter中，然后再按照最常见的顺序写到文件中，0号位置放置用于填充的词，1号位置放置未曾出现在词表的，这种顺序将常见的字在空间中聚集起来，私以为对模型学习有所帮助。得到vocab之后，就是获取vocab与id之间的映射关系了：

def load_vocab(vocab_file):
    vocab = open(vocab_file, encoding='utf8').read().splitlines()
    vocab2id = {x:i for i,x in enumerate(vocab)}
    return vocab2id

vocab2id = load_vocab('vocab.txt')

接着就是加载数据了，我将数据处理成列表，每一个元素为(文本,是否存在实体)

def exists(obj):
    """目标是否存在"""
    if obj:
        return True
    else:
        return False

def load_data():
    """加载数据"""
    out = []
    with open(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/labeled_data_for_mrc.jsonl', 'r', encoding='utf8') as f:
        for i, line in enumerate(f):
            if i == 1000:
                break
            line = json.loads(line)
            text = re.sub('\s+', '', line['text'])
            out.append([text, int(exists(line['labels']))])

    updated_files = glob.glob(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/updated*')
    for file in updated_files:
        with open(file, 'r', encoding='utf8') as f:
            for i, line in enumerate(f):
                line = json.loads(line)
                text = re.sub('\s+', '', line['text'])
                out.append([text, int(exists(line['label']))])

    print(f"total samples: {len(out)} has labels: {sum(list(zip(*out))[1])}")
    return out


def split_data(data):
    """切分数据集"""
    np.random.shuffle(data)
    length = len(data)
    train_data, valid_data, test_data = data[:int(length * 0.8)], data[int(length * 0.8):int(length * 0.9)], data[int(length * 0.9):]
    return train_data, valid_data, test_data


data = load_data()
train_data, valid_data, test_data = split_data(data)

#total samples: 1122 has labels: 492

我的数据一共包含1122条数据，其中有标签的有492条。接着采用torch的Dataset与DateLoader构建数据生成类，首先定义了tokenize函数，将文本映射成数字向量，定义了padding函数sequence_padding，collate_fn用于将单条数据合并成一个batch:

def tokenize(text):
    """将文本映射成数字向量"""
    return [vocab2id[x] if x in vocab2id else vocab2id[1] for x in text]

def sequence_padding(inputs, length=None, value=0, seq_dims=1, mode='post'):
    """Numpy函数，将序列padding到同一长度
    """
    if length is None:
        length = np.max([np.shape(x)[:seq_dims] for x in inputs], axis=0)
    elif not hasattr(length, '__getitem__'):
        length = [length]

    slices = [np.s_[:length[i]] for i in range(seq_dims)]
    slices = tuple(slices) if len(slices) > 1 else slices[0]
    pad_width = [(0, 0) for _ in np.shape(inputs[0])]

    outputs = []
    for x in inputs:
        x = x[slices]
        for i in range(seq_dims):
            if mode == 'post':
                pad_width[i] = (0, length[i] - np.shape(x)[i])
            elif mode == 'pre':
                pad_width[i] = (length[i] - np.shape(x)[i], 0)
            else:
                raise ValueError('"mode" argument must be "post" or "pre".')
        x = np.pad(x, pad_width, 'constant', constant_values=value)
        outputs.append(x)

    return np.array(outputs)

class BidDataset(Dataset):
    def __init__(self, data):
        super(BidDataset, self).__init__()
        self.data = data

    def __getitem__(self, index):
        d = self.data[index]  # 文本，是否存在标签
        input_ids = tokenize(d[0])
        labels = d[1]
        mask = [1] * len(input_ids)
        return input_ids, labels, mask

    def __len__(self):
        return len(self.data)


def collate_fn(batch):
    input_ids, labels, mask = list(zip(*batch))

    input_ids = torch.LongTensor(sequence_padding(input_ids))
    labels = torch.LongTensor(sequence_padding(labels))
    mask = torch.LongTensor(sequence_padding(mask))

    return input_ids, labels, mask

def get_dataloader(dataset):
    return DataLoader(dataset, batch_size=8, collate_fn=collate_fn)

train_dataset, valid_dataset, test_dataset = BidDataset(train_data), BidDataset(valid_data), BidDataset(test_data)
train_dataloader, valid_dataloader, test_dataloader = get_dataloader(train_dataset), get_dataloader(valid_dataset), get_dataloader(test_dataset)

数据处理完后，就可以定义模型了，模型我们采用DGCNN，其拥有较大的感知视野和较快的运行速度：

class ResidualGatedConv1D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation_rate):
        super(ResidualGatedConv1D, self).__init__()
        self.out_channels = out_channels
        self.conv1d = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels * 2,
            kernel_size=kernel_size,
            dilation=dilation_rate,
            padding=dilation_rate
        )
        self.layernorm = nn.LayerNorm([out_channels])
        self.alpha = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        x = x.transpose(2,1)
        outputs = self.conv1d(x)
        gate = torch.sigmoid(outputs[:, self.out_channels:])
        outputs = outputs[:, :self.out_channels] * gate
        outputs = self.layernorm(outputs.transpose(2,1))
        x = x.transpose(2,1) + self.alpha * outputs

        return x


class GlobalAveragePopl1D(nn.Module):
    """对某一维进行平均"""
    def __init__(self):
        super(GlobalAveragePopl1D, self).__init__()

    def forward(self, x):
        return torch.mean(x, dim=1)

class DGCNN(nn.Module):
    def __init__(self):
        super(DGCNN, self).__init__()
        self.dgcnn = nn.Sequential(
            nn.Embedding(len(vocab2id), 256, padding_idx=0),
            ResidualGatedConv1D(256, 256, 3, 1),
            nn.Dropout(0.1),
            ResidualGatedConv1D(256, 256, 3, 2),
            nn.Dropout(0.1),
            ResidualGatedConv1D(256, 256, 3, 4),
            nn.Dropout(0.1),
            ResidualGatedConv1D(256, 256, 3, 8),
            nn.Dropout(0.1),
            ResidualGatedConv1D(256, 256, 3, 1),
            nn.Dropout(0.1),
            ResidualGatedConv1D(256, 256, 3, 1),
            nn.Dropout(0.1),
            GlobalAveragePopl1D(),
            nn.Linear(256, 256),
            nn.Dropout(0.1),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.dgcnn(x)

定义好模型之后，就可以愉快的开始训练咯

def loss_fn(y_true, y_pred):
    loss = nn.BCELoss()(y_pred, y_true)
    return loss

def acc_metric(y_true, y_pred):
    y_pred = y_pred > 0.5
    correct = torch.sum(y_true == y_pred)
    return correct / y_true.shape[0]


def train():
    model = DGCNN()
    model.cuda()
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)
    best_acc = 0
    for _ in range(40):
        model.train()
        total_loss = 0
        total_acc = 0
        pbar = tqdm(enumerate(train_dataloader, 1), desc='train')
        for batch_id, batch in pbar:
            input_ids, label, mask = batch
            input_ids, label = input_ids.cuda(), label.cuda()
            logits = model(input_ids)
            loss = loss_fn(y_true=label, y_pred=logits)
            acc = acc_metric(y_true=label, y_pred=logits)
            total_loss += loss.item()
            total_acc += acc.item()

            pbar.set_postfix(loss=total_loss / batch_id, acc=total_acc / batch_id)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        pbar = tqdm(enumerate(valid_dataloader, 1), desc='dev')
        model.eval()
        total_acc = 0
        for batch_id, batch in pbar:
            input_ids, label, mask = batch
            input_ids, label = input_ids.cuda(), label.cuda()
            with torch.no_grad():
                logits = model(input_ids)
            acc = acc_metric(y_true=label, y_pred=logits)
            total_acc += acc.item()
            pbar.set_postfix(acc=total_acc / batch_id)

        if total_acc / batch_id > best_acc:
            best_acc = total_acc / batch_id
            torch.save(model.state_dict(), 'best_model.pt')
            print(f'best model saved at epoch {_} with best acc {best_acc}')


def evaluate():
    model = DGCNN()
    model.load_state_dict(torch.load('best_model.pt'))
    model.cuda()
    pbar = tqdm(enumerate(train_dataloader, 1), desc='dev')
    model.eval()
    total_acc = 0
    for batch_id, batch in pbar:
        input_ids, label, mask = batch
        input_ids, label = input_ids.cuda(), label.cuda()
        with torch.no_grad():
            logits = model(input_ids)
        acc = acc_metric(y_true=label, y_pred=logits)
        total_acc += acc.item()
        pbar.set_postfix(acc=total_acc / batch_id)

if __name__ == '__main__':
    train()
    evaluate()

此模型的准确率不错，可以达到90%左右。

改进

然而此精度并不能满足实际使用需求，因此进行了一些改进：

• 将DGCNN换成textCNN，尝试了多种参数的变化，包括核数目，卷积滤波器数量，dropout大小等，最终发现dropout=0.1，filters=100,kernel_sizes=[3,4,5]效果最佳
• 采用线性递增递减的学习率策略，尝试了多种10，20，30，40等epochs数，最大学习率等，最终发现0.0005的最大学习率，10的epochs数最好

通过这些改进，将准确率提升到了92%，满足实际使用需求了，整体改进代码如下：

import json
import re
import glob
import torch
import numpy as np
import torch.nn.functional as F
np.random.seed(42)
from collections import Counter
from torch.utils.data import Dataset, DataLoader
from torch import nn
from tqdm import tqdm
from transformers import AdamW, get_polynomial_decay_schedule_with_warmup

torch.manual_seed(3407)
torch.cuda.manual_seed(3407)
torch.cuda.manual_seed_all(3407)
HIDDEN_SIZE = 300
EPOCHS = 10

def load_vocab(vocab_file):
    vocab = open(vocab_file, encoding='utf8').read().splitlines()
    vocab2id = {x:i for i,x in enumerate(vocab)}
    return vocab, vocab2id

vocab, vocab2id = load_vocab('vocab.txt')

def exists(obj):
    """目标是否存在"""
    if obj:
        return True
    else:
        return False

def load_data():
    """加载数据"""
    out = []
    # with open(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/labeled_data_for_mrc.jsonl', 'r', encoding='utf8') as f:
    #     for i, line in enumerate(f):
    #         if i == 1000:
    #             break
    #         line = json.loads(line)
    #         text = re.sub('\s+', '', line['text'])
    #         out.append([text, int(exists(line['labels']))])
    #
    # updated_files = glob.glob(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/updated*')
    # for file in updated_files:
    #     with open(file, 'r', encoding='utf8') as f:
    #         for i, line in enumerate(f):
    #             line = json.loads(line)
    #             text = re.sub('\s+', '', line['text'])
    #             out.append([text, int(exists(line['label']))])

    with open(r'D:\PekingInfoOtherSearch\bert-mrc-pytorch\predicted_labeled_data_for_mrc.jsonl', 'r', encoding='utf8') as f:
        for i, line in enumerate(f):
            line = json.loads(line)
            text = re.sub('\s+', '', line['text'])
            out.append([text, int(exists(line['labels']))])

    print(f"total samples: {len(out)} has labels: {sum(list(zip(*out))[1])}")
    return out


def split_data(data):
    """切分数据集"""
    np.random.shuffle(data)
    length = len(data)
    train_data, valid_data, test_data = data[:int(length * 0.8)], data[int(length * 0.8):int(length * 0.9)], data[int(length * 0.9):]
    return train_data, valid_data, test_data


data = load_data()
train_data, valid_data, test_data = split_data(data)


def tokenize(text):
    """将文本映射成数字向量"""
    return [vocab2id[x] if x in vocab2id else vocab2id['[UNK]'] for x in text]

def sequence_padding(inputs, length=None, value=0, seq_dims=1, mode='post'):
    """Numpy函数，将序列padding到同一长度
    """
    if length is None:
        length = np.max([np.shape(x)[:seq_dims] for x in inputs], axis=0)
    elif not hasattr(length, '__getitem__'):
        length = [length]

    slices = [np.s_[:length[i]] for i in range(seq_dims)]
    slices = tuple(slices) if len(slices) > 1 else slices[0]
    pad_width = [(0, 0) for _ in np.shape(inputs[0])]

    outputs = []
    for x in inputs:
        x = x[slices]
        for i in range(seq_dims):
            if mode == 'post':
                pad_width[i] = (0, length[i] - np.shape(x)[i])
            elif mode == 'pre':
                pad_width[i] = (length[i] - np.shape(x)[i], 0)
            else:
                raise ValueError('"mode" argument must be "post" or "pre".')
        x = np.pad(x, pad_width, 'constant', constant_values=value)
        outputs.append(x)

    return np.array(outputs)

class BidDataset(Dataset):
    def __init__(self, data):
        super(BidDataset, self).__init__()
        self.data = data

    def __getitem__(self, index):
        d = self.data[index]  # 文本，是否存在标签
        input_ids = tokenize(d[0])
        labels = [d[1]]
        mask = [1] * len(input_ids)
        return input_ids, labels, mask

    def __len__(self):
        return len(self.data)


def collate_fn(batch):
    input_ids, labels, mask = list(zip(*batch))

    input_ids = torch.LongTensor(sequence_padding(input_ids))
    labels = torch.FloatTensor(sequence_padding(labels))
    mask = torch.LongTensor(sequence_padding(mask))

    return input_ids, labels, mask

def get_dataloader(dataset):
    return DataLoader(dataset, batch_size=8, collate_fn=collate_fn)

train_dataset, valid_dataset, test_dataset = BidDataset(train_data), BidDataset(valid_data), BidDataset(test_data)
train_dataloader, valid_dataloader, test_dataloader = get_dataloader(train_dataset), get_dataloader(valid_dataset), get_dataloader(test_dataset)


class ResidualGatedConv1D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation_rate):
        super(ResidualGatedConv1D, self).__init__()
        self.out_channels = out_channels
        self.conv1d = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels * 2,
            kernel_size=kernel_size,
            dilation=dilation_rate,
            padding=dilation_rate
        )
        self.layernorm = nn.LayerNorm([out_channels])
        self.alpha = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        x = x.transpose(2,1)
        outputs = self.conv1d(x)
        gate = torch.sigmoid(outputs[:, self.out_channels:])
        outputs = outputs[:, :self.out_channels] * gate
        outputs = self.layernorm(outputs.transpose(2,1))
        x = x.transpose(2,1) + self.alpha * outputs

        return x


class GlobalAveragePopl1D(nn.Module):
    """对某一维进行平均"""
    def __init__(self):
        super(GlobalAveragePopl1D, self).__init__()

    def forward(self, x):
        return torch.mean(x, dim=1)

class Embedding(nn.Module):
    def __init__(self):
        super(Embedding, self).__init__()
        self.embed = nn.Embedding(len(vocab2id), HIDDEN_SIZE, padding_idx=0)
        # self.embed.weight.data.copy_(torch.tensor(embedding).float())

    def forward(self, x):
        return self.embed(x)

class DGCNN(nn.Module):
    def __init__(self):
        super(DGCNN, self).__init__()
        drop_rate = 0.1
        hidden_size = 300
        self.dgcnn = nn.Sequential(
            Embedding(),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 1),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 2),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 4),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 8),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 1),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 1),
            nn.Dropout(drop_rate),
            GlobalAveragePopl1D(),
            nn.Linear(hidden_size, hidden_size),
            nn.Dropout(drop_rate),
            nn.Linear(hidden_size, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.dgcnn(x)


class textCNN(nn.Module):

    def __init__(self):
        super(textCNN, self).__init__()
        self.embed = Embedding()

        kernel_wins = [3,4,5]
        dim_channel = 100
        # Convolutional Layers with different window size kernels
        self.convs = nn.ModuleList([nn.Conv2d(1, dim_channel, (w, HIDDEN_SIZE)) for w in kernel_wins])
        # Dropout layer
        self.dropout = nn.Dropout(0.1)

        # FC layer
        self.fc = nn.Linear(len(kernel_wins) * dim_channel, 1)

    def forward(self, x):
        emb_x = self.embed(x)
        emb_x = emb_x.unsqueeze(1)

        con_x = [conv(emb_x) for conv in self.convs]

        pool_x = [F.max_pool1d(x.squeeze(-1), x.size()[2]) for x in con_x]

        fc_x = torch.cat(pool_x, dim=1)

        fc_x = fc_x.squeeze(-1)

        fc_x = self.dropout(fc_x)
        logit = torch.sigmoid(self.fc(fc_x))
        return logit

def loss_fn(y_true, y_pred):
    loss = nn.BCELoss()(y_pred, y_true)
    return loss

def acc_metric(y_true, y_pred):
    y_pred = (y_pred > 0.5).float()
    correct = torch.sum(y_true == y_pred)
    acc = correct / y_true.shape[0]
    recall = torch.sum(y_true * y_pred) / torch.sum(y_true).clamp(1e-9)
    precision = torch.sum(y_true * y_pred) / torch.sum(y_pred).clamp(1e-9)
    return acc, recall, precision


def build_optimizer_and_scheduler(model, warmup_proportion, total_steps):

    module = (model.module if hasattr(model, "module") else model)
    model_param = module.parameters()

    warmup_steps = int(warmup_proportion * total_steps)
    optimizer = AdamW(model_param, lr=0.0005, eps=1e-8)
    scheduler = get_polynomial_decay_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, lr_end=1e-5)
    return optimizer, scheduler


def train():
    model = textCNN()
    model.cuda()
    optimizer, scheduler = build_optimizer_and_scheduler(model, 0.1, len(train_dataloader)*EPOCHS)
    best_acc = 0
    for _ in range(EPOCHS):
        model.train()
        total_loss = 0
        total_acc = 0
        total_recall = 0
        total_precison = 0
        pbar = tqdm(enumerate(train_dataloader, 1), desc='train')
        for batch_id, batch in pbar:
            input_ids, label, mask = batch
            input_ids, label = input_ids.cuda(), label.cuda()
            logits = model(input_ids)
            loss = loss_fn(y_true=label, y_pred=logits)
            acc, recall, precision = acc_metric(y_true=label, y_pred=logits)
            total_loss += loss.item()
            total_acc += acc.item()
            total_recall += recall.item()
            total_precison += precision.item()

            pbar.set_description(f'Epoch {_}/{EPOCHS}')
            pbar.set_postfix(loss=total_loss / batch_id,
                             acc=total_acc / batch_id,
                             recall = total_recall / batch_id,
                             precision = total_precison / batch_id,
                             lr=optimizer.param_groups[0]["lr"])

            torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
            loss.backward()
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

        pbar = tqdm(enumerate(valid_dataloader, 1), desc='dev')
        model.eval()
        total_acc = 0
        total_recall = 0
        total_precison = 0
        for batch_id, batch in pbar:
            input_ids, label, mask = batch
            input_ids, label = input_ids.cuda(), label.cuda()
            with torch.no_grad():
                logits = model(input_ids)
            acc, recall, precision = acc_metric(y_true=label, y_pred=logits)
            total_acc += acc.item()
            total_recall += recall.item()
            total_precison += precision.item()
            pbar.set_postfix(acc=total_acc / batch_id,
                             recall=total_recall / batch_id,
                             precision=total_precison / batch_id,
                             )

        if total_acc / batch_id > best_acc:
            best_acc = total_acc / batch_id
            torch.save(model.state_dict(), 'best_model.pt')
            print(f'best model saved at epoch {_} with best acc {best_acc}')


def evaluate():
    model = DGCNN()
    model.load_state_dict(torch.load('best_model.pt'))
    model.cuda()
    pbar = tqdm(enumerate(test_dataloader, 1), desc='test')
    model.eval()
    total_acc = 0
    total_recall = 0
    total_precison = 0
    for batch_id, batch in pbar:
        input_ids, label, mask = batch
        input_ids, label = input_ids.cuda(), label.cuda()
        with torch.no_grad():
            logits = model(input_ids)
        acc, recall, precision = acc_metric(y_true=label, y_pred=logits)
        total_acc += acc.item()
        total_recall += recall.item()
        total_precison += precision.item()
        pbar.set_postfix(acc=total_acc / batch_id,
                         recall=total_recall / batch_id,
                         precision=total_precison / batch_id,
                         )

def convert2onnx():
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    import torch
    if torch.cuda.is_available():
        device = 'cuda:0'
    else:
        device = 'cpu'

    model = DGCNN()
    model.load_state_dict(torch.load('best_model.pt', map_location=device))
    model.to(device)
    model.eval()

    x = torch.zeros(1, 300, requires_grad=True).long()

    torch.onnx.export(model,               # model being run
                      x,                         # model input (or a tuple for multiple inputs)
                      "best_model.onnx",   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=14,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      input_names = ['x'],   # the model's input names
                      output_names = ['output'], # the model's output names
                      dynamic_axes={'x' : {0 : 'batch_size', 1: 'seqlen'},# variable length axes
                                    'output' : {0 : 'batch_size', 1: 'seqlen'}})

if __name__ == '__main__':
    train()
    # evaluate()
    # convert2onnx()

继续改进

上文提到预训练好的词向量，因此我也比较了使用jieba分词核sougou预训练词向量结合的模型效果，最好的结果为准确率0.927，提升了0.7个百分点，也算是不错的提升，最终在实际使用中，采用词向量模型，整体代码如下：

import json
import re
import glob
import torch
import jieba
import numpy as np
import torch.nn.functional as F
np.random.seed(42)
from collections import Counter
from torch.utils.data import Dataset, DataLoader
from torch import nn
from tqdm import tqdm
from transformers import AdamW, get_polynomial_decay_schedule_with_warmup

torch.manual_seed(3407)
torch.cuda.manual_seed(3407)
torch.cuda.manual_seed_all(3407)
HIDDEN_SIZE = 300
EPOCHS = 10

def load_vocab(vocab_file):
    vocab = open(vocab_file, encoding='utf8').read().splitlines()
    vocab2id = {x:i for i,x in enumerate(vocab)}
    return vocab, vocab2id

vocab, vocab2id = load_vocab('word_vocab.txt')


def load_embedding():
    vocab2embed = {}
    with open(r'D:\PekingInfoResearch\pretrain_models\word2vec\sgns.sogou.char', encoding='utf8') as f:
        f.readline()
        for line in tqdm(f, 'load embedding'):
            line = line.split()
            vocab2embed[line[0]] = list(map(float, line[1:]))
    out_embedding = []
    for word in vocab:
        if word in vocab2embed:
            out_embedding.append(vocab2embed[word])
        else:
            out_embedding.append(np.zeros(300))
    return np.array(out_embedding)


def exists(obj):
    """目标是否存在"""
    if obj:
        return True
    else:
        return False

def load_data():
    """加载数据"""
    out = []
    # with open(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/labeled_data_for_mrc.jsonl', 'r', encoding='utf8') as f:
    #     for i, line in enumerate(f):
    #         if i == 1000:
    #             break
    #         line = json.loads(line)
    #         text = re.sub('\s+', '', line['text'])
    #         out.append([text, int(exists(line['labels']))])
    #
    # updated_files = glob.glob(r'D:\open_data\ner\bid_data\labeled_data_for_mrc/updated*')
    # for file in updated_files:
    #     with open(file, 'r', encoding='utf8') as f:
    #         for i, line in enumerate(f):
    #             line = json.loads(line)
    #             text = re.sub('\s+', '', line['text'])
    #             out.append([text, int(exists(line['label']))])

    with open(r'D:\PekingInfoOtherSearch\bert-mrc-pytorch\predicted_labeled_data_for_mrc.jsonl', 'r', encoding='utf8') as f:
        for i, line in enumerate(f):
            line = json.loads(line)
            text = re.sub('\s+', '', line['text'])
            out.append([text, int(exists(line['labels']))])

    print(f"total samples: {len(out)} has labels: {sum(list(zip(*out))[1])}")
    return out


def split_data(data):
    """切分数据集"""
    np.random.shuffle(data)
    length = len(data)
    train_data, valid_data, test_data = data[:int(length * 0.8)], data[int(length * 0.8):int(length * 0.9)], data[int(length * 0.9):]
    return train_data, valid_data, test_data


data = load_data()
train_data, valid_data, test_data = split_data(data)


def tokenize(text):
    """将文本映射成数字向量"""
    return [vocab2id[x] if x in vocab2id else vocab2id['[UNK]'] for x in text]

def sequence_padding(inputs, length=None, value=0, seq_dims=1, mode='post'):
    """Numpy函数，将序列padding到同一长度
    """
    if length is None:
        length = np.max([np.shape(x)[:seq_dims] for x in inputs], axis=0)
    elif not hasattr(length, '__getitem__'):
        length = [length]

    slices = [np.s_[:length[i]] for i in range(seq_dims)]
    slices = tuple(slices) if len(slices) > 1 else slices[0]
    pad_width = [(0, 0) for _ in np.shape(inputs[0])]

    outputs = []
    for x in inputs:
        x = x[slices]
        for i in range(seq_dims):
            if mode == 'post':
                pad_width[i] = (0, length[i] - np.shape(x)[i])
            elif mode == 'pre':
                pad_width[i] = (length[i] - np.shape(x)[i], 0)
            else:
                raise ValueError('"mode" argument must be "post" or "pre".')
        x = np.pad(x, pad_width, 'constant', constant_values=value)
        outputs.append(x)

    return np.array(outputs)

class BidDataset(Dataset):
    def __init__(self, data):
        super(BidDataset, self).__init__()
        self.data = data

    def __getitem__(self, index):
        d = self.data[index]  # 文本，是否存在标签
        input_ids = tokenize(jieba.lcut(d[0]))
        labels = [d[1]]
        mask = [1] * len(input_ids)
        return input_ids, labels, mask

    def __len__(self):
        return len(self.data)


def collate_fn(batch):
    input_ids, labels, mask = list(zip(*batch))

    input_ids = torch.LongTensor(sequence_padding(input_ids))
    labels = torch.FloatTensor(sequence_padding(labels))
    mask = torch.LongTensor(sequence_padding(mask))

    return input_ids, labels, mask

def get_dataloader(dataset):
    return DataLoader(dataset, batch_size=8, collate_fn=collate_fn)

train_dataset, valid_dataset, test_dataset = BidDataset(train_data), BidDataset(valid_data), BidDataset(test_data)
train_dataloader, valid_dataloader, test_dataloader = get_dataloader(train_dataset), get_dataloader(valid_dataset), get_dataloader(test_dataset)


class ResidualGatedConv1D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation_rate):
        super(ResidualGatedConv1D, self).__init__()
        self.out_channels = out_channels
        self.conv1d = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels * 2,
            kernel_size=kernel_size,
            dilation=dilation_rate,
            padding=dilation_rate
        )
        self.layernorm = nn.LayerNorm([out_channels])
        self.alpha = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        x = x.transpose(2,1)
        outputs = self.conv1d(x)
        gate = torch.sigmoid(outputs[:, self.out_channels:])
        outputs = outputs[:, :self.out_channels] * gate
        outputs = self.layernorm(outputs.transpose(2,1))
        x = x.transpose(2,1) + self.alpha * outputs

        return x


class GlobalAveragePopl1D(nn.Module):
    """对某一维进行平均"""
    def __init__(self):
        super(GlobalAveragePopl1D, self).__init__()

    def forward(self, x):
        return torch.mean(x, dim=1)

class Embedding(nn.Module):
    def __init__(self, embedding=None):
        super(Embedding, self).__init__()
        self.embed = nn.Embedding(len(vocab2id), HIDDEN_SIZE, padding_idx=0)
        if embedding:
            self.embed.weight.data.copy_(torch.tensor(embedding).float())

    def forward(self, x):
        return self.embed(x)

class DGCNN(nn.Module):
    def __init__(self):
        super(DGCNN, self).__init__()
        drop_rate = 0.1
        hidden_size = 300
        self.dgcnn = nn.Sequential(
            Embedding(),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 1),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 2),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 4),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 8),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 1),
            nn.Dropout(drop_rate),
            ResidualGatedConv1D(hidden_size, hidden_size, 3, 1),
            nn.Dropout(drop_rate),
            GlobalAveragePopl1D(),
            nn.Linear(hidden_size, hidden_size),
            nn.Dropout(drop_rate),
            nn.Linear(hidden_size, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.dgcnn(x)


class textCNN(nn.Module):

    def __init__(self, embedding=None):
        super(textCNN, self).__init__()
        self.embed = Embedding(embedding)

        kernel_wins = [3,4,5]
        dim_channel = 100
        # Convolutional Layers with different window size kernels
        self.convs = nn.ModuleList([nn.Conv2d(1, dim_channel, (w, HIDDEN_SIZE)) for w in kernel_wins])
        # Dropout layer
        self.dropout = nn.Dropout(0.1)

        # FC layer
        self.fc = nn.Linear(len(kernel_wins) * dim_channel, 1)

    def forward(self, x):
        emb_x = self.embed(x)
        emb_x = emb_x.unsqueeze(1)

        con_x = [conv(emb_x) for conv in self.convs]

        pool_x = [F.adaptive_max_pool1d(x.squeeze(-1), 1) for x in con_x]

        fc_x = torch.cat(pool_x, dim=1)

        fc_x = fc_x.squeeze(-1)

        fc_x = self.dropout(fc_x)
        logit = torch.sigmoid(self.fc(fc_x))
        return logit

def loss_fn(y_true, y_pred):
    loss = nn.BCELoss()(y_pred, y_true)
    return loss

def acc_metric(y_true, y_pred):
    y_pred = (y_pred > 0.5).float()
    correct = torch.sum(y_true == y_pred)
    acc = correct / y_true.shape[0]
    recall = torch.sum(y_true * y_pred) / torch.sum(y_true).clamp(1e-9)
    precision = torch.sum(y_true * y_pred) / torch.sum(y_pred).clamp(1e-9)
    return acc, recall, precision


def build_optimizer_and_scheduler(model, warmup_proportion, total_steps):

    module = (model.module if hasattr(model, "module") else model)
    model_param = module.parameters()

    warmup_steps = int(warmup_proportion * total_steps)
    optimizer = AdamW(model_param, lr=0.001, eps=1e-8)
    scheduler = get_polynomial_decay_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, lr_end=1e-5)
    return optimizer, scheduler


def train(embedding=None):
    model = textCNN(embedding)
    model.cuda()
    optimizer, scheduler = build_optimizer_and_scheduler(model, 0.1, len(train_dataloader)*EPOCHS)
    best_acc = 0
    for _ in range(EPOCHS):
        model.train()
        total_loss = 0
        total_acc = 0
        total_recall = 0
        total_precison = 0
        pbar = tqdm(enumerate(train_dataloader, 1), desc='train')
        for batch_id, batch in pbar:
            input_ids, label, mask = batch
            input_ids, label = input_ids.cuda(), label.cuda()
            logits = model(input_ids)
            loss = loss_fn(y_true=label, y_pred=logits)
            acc, recall, precision = acc_metric(y_true=label, y_pred=logits)
            total_loss += loss.item()
            total_acc += acc.item()
            total_recall += recall.item()
            total_precison += precision.item()

            pbar.set_description(f'Epoch {_}/{EPOCHS}')
            pbar.set_postfix(loss=total_loss / batch_id,
                             acc=total_acc / batch_id,
                             recall = total_recall / batch_id,
                             precision = total_precison / batch_id,
                             lr=optimizer.param_groups[0]["lr"])

            torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
            loss.backward()
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

        pbar = tqdm(enumerate(valid_dataloader, 1), desc='dev')
        model.eval()
        total_acc = 0
        total_recall = 0
        total_precison = 0
        for batch_id, batch in pbar:
            input_ids, label, mask = batch
            input_ids, label = input_ids.cuda(), label.cuda()
            with torch.no_grad():
                logits = model(input_ids)
            acc, recall, precision = acc_metric(y_true=label, y_pred=logits)
            total_acc += acc.item()
            total_recall += recall.item()
            total_precison += precision.item()
            pbar.set_postfix(acc=total_acc / batch_id,
                             recall=total_recall / batch_id,
                             precision=total_precison / batch_id,
                             )

        if total_acc / batch_id > best_acc:
            best_acc = total_acc / batch_id
            torch.save(model.state_dict(), 'best_model.word.pt')
            print(f'best model saved at epoch {_} with best acc {best_acc}')


def evaluate():
    model = textCNN()
    model.load_state_dict(torch.load('best_model.word.pt'))
    model.cuda()
    pbar = tqdm(enumerate(test_dataloader, 1), desc='test')
    model.eval()
    total_acc = 0
    total_recall = 0
    total_precison = 0
    for batch_id, batch in pbar:
        input_ids, label, mask = batch
        input_ids, label = input_ids.cuda(), label.cuda()
        with torch.no_grad():
            logits = model(input_ids)
        acc, recall, precision = acc_metric(y_true=label, y_pred=logits)
        total_acc += acc.item()
        total_recall += recall.item()
        total_precison += precision.item()
        pbar.set_postfix(acc=total_acc / batch_id,
                         recall=total_recall / batch_id,
                         precision=total_precison / batch_id,
                         )

def convert2onnx():
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    import torch
    if torch.cuda.is_available():
        device = 'cuda:0'
    else:
        device = 'cpu'

    model = textCNN()
    model.load_state_dict(torch.load('best_model.word.pt', map_location=device))
    model.to(device)
    model.eval()

    x = torch.zeros(1, 300, requires_grad=True).long()

    torch.onnx.export(model,               # model being run
                      x,                         # model input (or a tuple for multiple inputs)
                      "best_model.word.onnx",   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=14,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      input_names = ['x'],   # the model's input names
                      output_names = ['output'], # the model's output names
                      dynamic_axes={'x' : {0 : 'batch_size', 1: 'seqlen'},# variable length axes
                                    'output' : {0 : 'batch_size', 1: 'seqlen'}})

if __name__ == '__main__':
    # embedding = load_embedding()
    # train()
    # evaluate()
    convert2onnx()

实际使用情况

实际使用情况就是，速度很快，效果也很不错！