-
Notifications
You must be signed in to change notification settings - Fork 209
/
Copy pathtrain_albert.py
153 lines (124 loc) · 5.02 KB
/
train_albert.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
# -*- coding:utf-8 -*-
'''
-------------------------------------------------
Description : albert train
Author : machinelp
Date : 2020-06-04
-------------------------------------------------
'''
import numpy as np
import pandas as pd
from keras.layers import *
from bert4keras.backend import keras, set_gelu
from bert4keras.bert import build_bert_model
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding, DataGenerator
from bert4keras.tokenizer import Tokenizer
import pandas as pd
import numpy as np
from textmatch.config.constant import Constant as const
from textmatch.models.text_embedding.albert_embedding import ALBertEmbedding
set_gelu('tanh') # 切换gelu版本
maxlen = 32
batch_size = 16
num_classes = 2
epochs = 20
learning_rate = 2e-5
# sim roeberta_zh
# 【百度网盘】链接:https://pan.baidu.com/s/1-PRsjQSwkGSpQkmjjnSXmw 密码:ynjs
config_path = 'albert_tiny_google_zh_489k/albert_config.json'
checkpoint_path = 'albert_tiny_google_zh_489k/albert_model.ckpt'
dict_path = 'albert_tiny_google_zh_489k/vocab.txt'
def load_data(filename):
D = []
data = pd.read_csv(filename)
data.dropna(axis=0, how='any', inplace=True)
data = data.values.tolist()
for per_data in data:
D.append( (per_data[0],per_data[1],int(per_data[2])) )
return D
# 加载数据集
train_val_data = load_data('./data/train_data.csv')
# test_data = load_data('dev.csv')
# 查看一下数据
print ( 'train>>>>', train_val_data[0] )
print ( '训练验证集数量:', len(train_val_data) )
random_order = range(len(train_val_data))
np.random.shuffle(list(random_order))
train_data = [train_val_data[j] for i, j in enumerate(random_order) if i % 5 != 1 ]
valid_data = [train_val_data[j] for i, j in enumerate(random_order) if i % 5 == 1 ]
test_data = valid_data
print ( '训练集数量:', len(train_data) )
print ( '验证集数量:', len(valid_data) )
print ( '测试集数量:', len(test_data) )
# 建立分词器
tokenizer = Tokenizer(dict_path, do_lower_case=True)
class data_generator(DataGenerator):
"""数据生成器
"""
def __iter__(self, random=False):
idxs = list(range(len(self.data)))
if random:
np.random.shuffle(idxs)
batch_token_ids, batch_segment_ids, batch_labels = [], [], []
for i in idxs:
text1, text2, label = self.data[i]
# print(text1, text2, label)
token_ids, segment_ids = tokenizer.encode(text1, text2, max_length=maxlen)
batch_token_ids.append(token_ids)
batch_segment_ids.append(segment_ids)
batch_labels.append([label])
if len(batch_token_ids) == self.batch_size or i == idxs[-1]:
batch_token_ids = sequence_padding(batch_token_ids)
batch_segment_ids = sequence_padding(batch_segment_ids)
batch_labels = sequence_padding(batch_labels)
yield [batch_token_ids, batch_segment_ids], batch_labels
batch_token_ids, batch_segment_ids, batch_labels = [], [], []
const.ALBERT_CONFIG_PATH = config_path
const.ALBERT_CHECKPOINT_PATH = checkpoint_path
const.ALBERT_DICT_PATH = dict_path
# 加载预训练模型
bert_embedding = ALBertEmbedding(const.ALBERT_CONFIG_PATH, const.ALBERT_CHECKPOINT_PATH, const.ALBERT_DICT_PATH, train_mode=True)
bert = bert_embedding.bert
output = Dropout(rate=0.1)(bert.model.output)
output = Dense(units=num_classes,
activation='softmax',
kernel_initializer=bert.initializer)(output)
model = keras.models.Model(bert.model.input, output)
model.summary()
model.compile(
loss='sparse_categorical_crossentropy',
optimizer=Adam(learning_rate), # 用足够小的学习率
# optimizer=PiecewiseLinearLearningRate(Adam(5e-5), {10000: 1, 30000: 0.1}),
metrics=['accuracy'],
)
# 转换数据集
train_generator = data_generator(train_data, batch_size)
valid_generator = data_generator(valid_data, batch_size)
test_generator = data_generator(test_data, batch_size)
def evaluate(data):
total, right = 0., 0.
for x_true, y_true in data:
y_pred = model.predict(x_true).argmax(axis=1)
y_true = y_true[:, 0]
total += len(y_true)
right += (y_true == y_pred).sum()
return right / total
class Evaluator(keras.callbacks.Callback):
def __init__(self):
self.best_val_acc = 0.
def on_epoch_end(self, epoch, logs=None):
val_acc = evaluate(valid_generator)
if val_acc > self.best_val_acc:
self.best_val_acc = val_acc
model.save_weights('best_model.weights')
test_acc = evaluate(test_generator)
print(u'val_acc: %.5f, best_val_acc: %.5f, test_acc: %.5f\n'
% (val_acc, self.best_val_acc, test_acc))
evaluator = Evaluator()
model.fit_generator(train_generator.forfit(),
steps_per_epoch=len(train_generator),
epochs=epochs,
callbacks=[evaluator])
model.load_weights('best_model.weights')
print(u'final test acc: %05f\n' % (evaluate(test_generator)))