-
Notifications
You must be signed in to change notification settings - Fork 11
/
Copy pathrun_classifier.py
239 lines (203 loc) · 10.5 KB
/
run_classifier.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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
# -*- coding: UTF-8 -*-
__author__ = 'huanghui'
__date__ = '2021/4/18 15:06'
__project__ = 'tfbert'
import json
import os
import argparse
import tensorflow.compat.v1 as tf
from tfbert import (
Trainer, Dataset,
SequenceClassification,
CONFIGS, TOKENIZERS, devices, set_seed)
from tfbert.data.classification import convert_examples_to_features, InputExample
from sklearn.metrics import accuracy_score
import pandas as pd
from typing import Dict
import numpy as np
def create_args():
parser = argparse.ArgumentParser()
parser.add_argument('--model_type', default='bert', type=str, choices=CONFIGS.keys())
parser.add_argument('--optimizer_type', default='adamw', type=str, help="优化器类型")
parser.add_argument('--model_dir', default='model_path', type=str,
help="预训练模型存放文件夹，文件夹下ckpt文件名为model.ckpt，"
"config文件名为config.json，词典文件名为vocab.txt")
parser.add_argument('--config_path', default=None, type=str, help="若配置文件名不是默认的，可在这里输入")
parser.add_argument('--vocab_path', default=None, type=str, help="若词典文件名不是默认的，可在这里输入")
parser.add_argument('--pretrained_checkpoint_path', default=None, type=str, help="若模型文件名不是默认的，可在这里输入")
parser.add_argument('--output_dir', default='output/classification', type=str, help="")
parser.add_argument('--export_dir', default='output/classification/pb', type=str, help="")
parser.add_argument('--labels', default='体育,娱乐,家居,房产,教育', type=str, help="文本分类标签")
parser.add_argument('--train_file', default='data/classification/train.csv', type=str, help="")
parser.add_argument('--dev_file', default='data/classification/dev.csv', type=str, help="")
parser.add_argument('--test_file', default='data/classification/test.csv', type=str, help="")
parser.add_argument("--num_train_epochs", default=3, type=int, help="训练轮次")
parser.add_argument("--max_seq_length", default=32, type=int, help="最大句子长度")
parser.add_argument("--batch_size", default=32, type=int, help="训练批次")
parser.add_argument("--gradient_accumulation_steps", default=1, type=int, help="梯度累积")
parser.add_argument("--learning_rate", default=2e-5, type=float, help="学习率")
parser.add_argument("--warmup_proportion", default=0.1, type=float,
help="Proportion of training to perform linear learning rate warmup for.")
parser.add_argument("--weight_decay", default=0.01, type=float, help="Weight decay if we apply some.")
parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")
parser.add_argument("--do_predict", action="store_true", help="Whether to run test on the test set.")
parser.add_argument("--evaluate_during_training", action="store_true", help="是否边训练边验证")
parser.add_argument("--do_export", action="store_true", help="将模型导出为pb格式.")
parser.add_argument("--logging_steps", default=1000, type=int, help="训练时每隔几步验证一次")
parser.add_argument("--saving_steps", default=1000, type=int, help="训练时每隔几步保存一次")
parser.add_argument("--random_seed", default=42, type=int, help="随机种子")
parser.add_argument("--threads", default=8, type=int, help="数据处理进程数")
parser.add_argument("--max_checkpoints", default=1, type=int, help="模型保存最大数量，默认只保存一个")
parser.add_argument("--single_device", action="store_true", help="是否只使用一个device，默认使用所有的device训练")
parser.add_argument("--use_xla", action="store_true", help="是否使用XLA加速")
parser.add_argument(
"--mixed_precision", action="store_true",
help="混合精度训练，tf下测试需要同时使用xla才有加速效果，但是开始编译很慢")
args = parser.parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
if not args.single_device:
args.batch_size = args.batch_size * len(devices())
args.labels = args.labels.split(',')
return args
def create_dataset(set_type, tokenizer, args):
filename_map = {
'train': args.train_file, 'dev': args.dev_file, 'test': args.test_file
}
examples = []
datas = pd.read_csv(filename_map[set_type], encoding='utf-8', sep='\t').values.tolist()
for data in datas:
examples.append(InputExample(
guid=0, text_a=data[1], label=data[0]
))
features = convert_examples_to_features(
examples, tokenizer,
max_length=args.max_seq_length, set_type=set_type,
label_list=args.labels, threads=args.threads)
dataset = Dataset(features,
is_training=bool(set_type == 'train'),
batch_size=args.batch_size,
drop_last=bool(set_type == 'train'),
buffer_size=len(features),
max_length=args.max_seq_length)
columns = ['input_ids', 'attention_mask', 'token_type_ids', 'label_ids']
if "pinyin_ids" in features[0] and features[0]['pinyin_ids'] is not None:
columns = ['input_ids', 'attention_mask', 'token_type_ids', 'pinyin_ids', 'label_ids']
dataset.format_as(columns)
return dataset
def get_model_fn(config, args):
def model_fn(inputs, is_training):
model = SequenceClassification(
model_type=args.model_type, config=config,
num_classes=len(args.labels), is_training=is_training,
**inputs)
outputs = {'outputs': {'logits': model.logits, 'label_ids': inputs['label_ids']}}
if model.loss is not None:
loss = model.loss / args.gradient_accumulation_steps
outputs['loss'] = loss
return outputs
return model_fn
def get_serving_fn(config, args):
def serving_fn():
input_ids = tf.placeholder(shape=[None, args.max_seq_length], dtype=tf.int64, name='input_ids')
attention_mask = tf.placeholder(shape=[None, args.max_seq_length], dtype=tf.int64, name='attention_mask')
token_type_ids = tf.placeholder(shape=[None, args.max_seq_length], dtype=tf.int64, name='token_type_ids')
if args.model_type == 'glyce_bert':
pinyin_ids = tf.placeholder(shape=[None, args.max_seq_length, 8], dtype=tf.int64, name='pinyin_ids')
else:
pinyin_ids = None
model = SequenceClassification(
model_type=args.model_type, config=config,
num_classes=len(args.labels), is_training=False,
input_ids=input_ids,
pinyin_ids=pinyin_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids
)
inputs = {'input_ids': input_ids, 'attention_mask': attention_mask, 'token_type_ids': token_type_ids}
if pinyin_ids is not None:
inputs['pinyin_ids'] = pinyin_ids
outputs = {'logits': model.logits}
return inputs, outputs
return serving_fn
def metric_fn(outputs: Dict) -> Dict:
"""
这里定义评估函数
:param outputs: trainer evaluate 返回的预测结果，model fn的outputs包含哪些字段就会有哪些字段
:return: 需要返回字典结果
"""
predictions = np.argmax(outputs['logits'], -1)
score = accuracy_score(outputs['label_ids'], predictions)
return {'accuracy': score}
def main():
args = create_args()
set_seed(args.random_seed)
config = CONFIGS[args.model_type].from_pretrained(
args.model_dir if args.config_path is None else args.config_path)
tokenizer = TOKENIZERS[args.model_type].from_pretrained(
args.model_dir if args.vocab_path is None else args.vocab_path, do_lower_case=True)
train_dataset, dev_dataset, predict_dataset = None, None, None
if args.do_train:
train_dataset = create_dataset('train', tokenizer, args)
if args.do_eval:
dev_dataset = create_dataset('dev', tokenizer, args)
if args.do_predict:
predict_dataset = create_dataset('test', tokenizer, args)
output_types, output_shapes = (train_dataset or dev_dataset or predict_dataset).output_types_and_shapes()
trainer = Trainer(
train_dataset=train_dataset,
eval_dataset=dev_dataset,
output_types=output_types,
output_shapes=output_shapes,
metric_fn=metric_fn,
use_xla=args.use_xla,
optimizer_type=args.optimizer_type,
learning_rate=args.learning_rate,
num_train_epochs=args.num_train_epochs,
gradient_accumulation_steps=args.gradient_accumulation_steps,
max_checkpoints=1,
max_grad=1.0,
warmup_proportion=args.warmup_proportion,
mixed_precision=args.mixed_precision,
single_device=args.single_device,
logging=True
)
trainer.build_model(model_fn=get_model_fn(config, args))
if args.do_train and train_dataset is not None:
# 训练阶段需要先compile优化器才能初始化权重
# 因为adam也是具备参数的
trainer.compile()
trainer.from_pretrained(
args.mo
5334
del_dir if args.pretrained_checkpoint_path is None else args.pretrained_checkpoint_path)
if args.do_train and train_dataset is not None:
trainer.train(
output_dir=args.output_dir,
evaluate_during_training=args.evaluate_during_training,
logging_steps=args.logging_steps,
saving_steps=args.saving_steps,
greater_is_better=True,
load_best_model=True,
metric_for_best_model='accuracy')
config.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir)
if args.do_eval and dev_dataset is not None:
eval_outputs = trainer.evaluate()
print(json.dumps(
eval_outputs, ensure_ascii=False, indent=4
))
if args.do_predict and predict_dataset is not None:
outputs = trainer.predict('test', ['logits'], dataset=predict_dataset)
label_ids = np.argmax(outputs['logits'], axis=-1)
labels = list(map(lambda x: args.labels[x], label_ids))
open(
os.path.join(args.output_dir, 'prediction.txt'), 'w', encoding='utf-8'
).write("\n".join(labels))
if args.do_export:
trainer.export(
get_serving_fn(config, args),
args.output_dir,
args.export_dir
)
if __name__ == '__main__':
main()