-
Notifications
You must be signed in to change notification settings - Fork 55
/
train_tener_cn.py
129 lines (110 loc) · 4.57 KB
/
train_tener_cn.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
from models.TENER import TENER
from fastNLP import cache_results
from fastNLP import Trainer, GradientClipCallback, WarmupCallback
from torch import optim
from fastNLP import SpanFPreRecMetric, BucketSampler
from fastNLP.embeddings import StaticEmbedding
from modules.pipe import CNNERPipe
import argparse
from modules.callbacks import EvaluateCallback
device = 0
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='resume', choices=['weibo', 'resume', 'ontonotes', 'msra'])
args = parser.parse_args()
dataset = args.dataset
if dataset == 'resume':
n_heads = 4
head_dims = 64
num_layers = 2
lr = 0.0007
attn_type = 'adatrans'
n_epochs = 50
elif dataset == 'weibo':
n_heads = 4
head_dims = 32
num_layers = 1
lr = 0.001
attn_type = 'adatrans'
n_epochs = 100
elif dataset == 'ontonotes':
n_heads = 4
head_dims = 48
num_layers = 2
lr = 0.0007
attn_type = 'adatrans'
n_epochs = 100
elif dataset == 'msra':
n_heads = 6
head_dims = 80
num_layers = 2
lr = 0.0007
attn_type = 'adatrans'
n_epochs = 100
pos_embed = None
batch_size = 16
warmup_steps = 0.01
after_norm = 1
model_type = 'transformer'
normalize_embed = True
dropout=0.15
fc_dropout=0.4
encoding_type = 'bmeso'
name = 'caches/{}_{}_{}_{}.pkl'.format(dataset, model_type, encoding_type, normalize_embed)
d_model = n_heads * head_dims
dim_feedforward = int(2 * d_model)
@cache_results(name, _refresh=False)
def load_data():
# 替换路径
if dataset == 'ontonotes':
paths = {'train':'../data/OntoNote4NER/train.char.bmes',
"dev":'../data/OntoNote4NER/dev.char.bmes',
"test":'../data/OntoNote4NER/test.char.bmes'}
min_freq = 2
elif dataset == 'weibo':
paths = {'train': '../data/WeiboNER/train.all.bmes',
'dev':'../data/WeiboNER/dev.all.bmes',
'test':'../data/WeiboNER/test.all.bmes'}
min_freq = 1
elif dataset == 'resume':
paths = {'train': '../data/ResumeNER/train.char.bmes',
'dev':'../data/ResumeNER/dev.char.bmes',
'test':'../data/ResumeNER/test.char.bmes'}
min_freq = 1
elif dataset == 'msra':
paths = {'train': '../data/MSRANER/train_dev.char.bmes',
'dev':'../data/MSRANER/test.char.bmes',
'test':'../data/MSRANER/test.char.bmes'}
min_freq = 2
data_bundle = CNNERPipe(bigrams=True, encoding_type=encoding_type).process_from_file(paths)
embed = StaticEmbedding(data_bundle.get_vocab('chars'),
model_dir_or_name='../data/gigaword_chn.all.a2b.uni.ite50.vec',
min_freq=1, only_norm_found_vector=normalize_embed, word_dropout=0.01, dropout=0.3)
bi_embed = StaticEmbedding(data_bundle.get_vocab('bigrams'),
model_dir_or_name='../data/gigaword_chn.all.a2b.bi.ite50.vec',
word_dropout=0.02, dropout=0.3, min_freq=min_freq,
only_norm_found_vector=normalize_embed, only_train_min_freq=True)
return data_bundle, embed, bi_embed
data_bundle, embed, bi_embed = load_data()
print(data_bundle)
model = TENER(tag_vocab=data_bundle.get_vocab('target'), embed=embed, num_layers=num_layers,
d_model=d_model, n_head=n_heads,
feedforward_dim=dim_feedforward, dropout=dropout,
after_norm=after_norm, attn_type=attn_type,
bi_embed=bi_embed,
fc_dropout=fc_dropout,
pos_embed=pos_embed,
scale=attn_type=='transformer')
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
callbacks = []
clip_callback = GradientClipCallback(clip_type='value', clip_value=5)
evaluate_callback = EvaluateCallback(data_bundle.get_dataset('test'))
if warmup_steps>0:
warmup_callback = WarmupCallback(warmup_steps, schedule='linear')
callbacks.append(warmup_callback)
callbacks.extend([clip_callback, evaluate_callback])
trainer = Trainer(data_bundle.get_dataset('train'), model, optimizer, batch_size=batch_size, sampler=BucketSampler(),
num_workers=2, n_epochs=n_epochs, dev_data=data_bundle.get_dataset('dev'),
metrics=SpanFPreRecMetric(tag_vocab=data_bundle.get_vocab('target'), encoding_type=encoding_type),
dev_batch_size=batch_size, callbacks=callbacks, device=device, test_use_tqdm=False,
use_tqdm=True, print_every=300, save_path=None)
trainer.train(load_best_model=False)