用Transformers的Trainer做Text-CNN文本分类

SUN_SU3

于 2022-01-12 10:22:12 发布

阅读量1.5k

点赞数 3

分类专栏：模型 NLP 文章标签：深度学习

本文链接：https://blog.youkuaiyun.com/u013546508/article/details/122447500

版权

本文介绍如何利用Transformers库中的Trainer类，借鉴其在Bert Fine-tuning上的便捷性，应用于Text-CNN模型进行文本分类任务，以简化训练过程并提升推理速度。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Transformers的Trainer这个类做Bert Fine-tuning时比较方便，近期用Text-CNN做文本分类实验提升推理速度时，就直接引用了Trainer的训练流程来训练text-CNN，简化训练步骤

import os
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from sklearn import metrics
from datasets import Dataset
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from transformers import Trainer, TrainingArguments, BertTokenizer

os.environ['CUDA_VISIBLE_DEVICES'] = '7'

# 加载数据
def create_dataset(data_file, tokenizer):
    print('create dataset')
    with open(data_file, 'r', encoding='utf-8') as f:
        data = [_.strip().split('\t') for _ in f.readlines()]
    x = [_[0] for _ in data]
    y = [int(_[1]) for _ in data]
    data_dict = {
   'text': x, 'label': y}
    dataset = Dataset.from_dict(data_dict)
    def

最低0.47元/天解锁文章