基于微调BERT模型的新闻分类

%reload_ext autoreload
%autoreload 2
%matplotlib inline
import os
# os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID";
# os.environ["CUDA_VISIBLE_DEVICES"]="0";   

采用ktrain库进行文本分类

目前ktrain工具箱已经集成transformers库，可以调用transformers库中的方法和函数

加载数据

从sklearn库中导入fetch_20newsgroups新闻数据库

共18000篇新闻文章，共20个不同类别

categories = ['alt.atheism', 'soc.religion.christian',
             'comp.graphics', 'sci.med']
from sklearn.datasets import fetch_20newsgroups
train_b = fetch_20newsgroups(subset='train',
   categories=categories, shuffle=True, random_state=42)
test_b = fetch_20newsgroups(subset='test',
   categories=categories, shuffle=True, random_state=42)

print('size of training set: %s' % (len(train_b['data'])))
print('size of validation set: %s' % (len(test_b['data'])))
print('classes: %s' % (train_b.target_names))

x_train = train_b.data
y_train = train_b.target
x_test = test_b.data
y_test = test_b.target
x_train

第一步：预处理数据并构造一个trasnformer模型

ktrain支持transformers库的各种内置模型(https://huggingface.co/transformers/pretrained_models.html) 和社区共享模型https://huggingface.co/models

import ktrain
from ktrain import text
MODEL_NAME = 'distilbert-base-uncased'
t = text.Transformer(MODEL_NAME, maxlen=500, class_names=train_b.target_names)
trn = t.preprocess_train(x_train, y_train)
val = t.preprocess_test(x_test, y_test)
model = t.get_classifier()
learner = ktrain.get_learner(model, train_data=trn, val_data=val, batch_size=6)

c

第二步[可选择]：估计一个较好的学习率

对于transformer模型来讲，学习率（learning rate）取2e-5 - 5e-5 会取得较好的性能. 这里，我们可以利用lr_find方法尝试在特定数据及寻找一个较好的学习率

这里为我们只运行2个epoch

learner.lr_find(show_plot=True, max_epochs=2)

第三步：训练模型

参数设置：learning_rate = 8e-5, epochs = 4

若在CPU上训练该模型，每个epoch大概耗时约1个小时…
因此推荐大家在带有高性能GPU的电脑上运行程序

learner.fit_onecycle(8e-5, 4)

第四步：评估模型

learner.validate(class_names=t.get_classes())

# 找到loss最大的样本
learner.view_top_losses(n=1, preproc=t)

# sci.med类型的新闻，但是都在探讨计算机图形学的内容--这种错误是可以理解的 
print(x_test[521])

第五步：在新的数据上进行测试

predictor = ktrain.get_predictor(learner.model, preproc=t)
predictor.predict('Jesus Christ is the central figure of Christianity.')

predictor.explain('Jesus Christ is the central figure of Christianity.')

predictor.save('/my_20newsgroup_predictor')
reloaded_predictor = ktrain.load_predictor('/my_20newsgroup_predictor')
reloaded_predictor.predict('Jesus Christ is the central figure of Christianity.')

reloaded_predictor.predict_proba('Jesus Christ is the central figure of Christianity.')

reloaded_predictor.get_classes()