# 模型训练脚本 # 链接：https://github.com/kangyishuai/NEWS-TEXT-CLASSIFICATION # 这里仅基于bert4torch实现了Top1解决方案中的finetune部分，直接使用了原作者的预训练权重转pytorch import numpy as np import pandas as pd from bert4torch.models import build_transformer_model, BaseModel from torch.utils.data import DataLoader from bert4torch.snippets import sequence_padding, ListDataset, Callback, EarlyStopping from bert4torch.tokenizers import Tokenizer import torch.nn.functional as F from sklearn.metrics import f1_score from sklearn.model_selection import StratifiedKFold import torch from torch import nn, optim from tqdm import tqdm # BERT base config_path = 'E:/Github/天池新闻分类/top1/pre_models/bert_config.json' checkpoint_path = 'E:/Github/天池新闻分类/top1/pre_models/pytorch_model.bin' dict_path = 'E:/Github/天池新闻分类/top1/pre_models/vocab.txt' device = f'cuda' if torch.cuda.is_available() else 'cpu' n = 5 # Cross-validation SEED = 2020 num_classes = 14 maxlen = 512 max_segment = 2 batch_size = 4 grad_accum_steps = 64 drop = 0.2 lr = 2e-5 epochs = 100 def load_data(df): """加载数据。""" D = list() for _, row in df.iterrows(): text = row['text'] label = row['label'] D.append((text, int(label))) return D # 建立分词器 tokenizer = Tokenizer(dict_path, do_lower_case=True) def sentence_split(words): """句子截断。""" document_len = len(words) index = list(range(0, document_len, maxlen-2)) index.append(document_len) segments = [] for i in range(len(index) - 1): segment = words[index[i]: index[i + 1]] assert len(segment) > 0 segment = tokenizer.tokens_to_ids(['[CLS]'] + segment + ['[SEP]']) segments.append(segment) assert len(segments) > 0 if len(segments) > max_segment: segment_ = int(max_segment / 2) return segments[:segment_] + segments[-segment_:] else: return segments def collate_fn(batch): batch_token_ids, batch_labels = [], [] for text, label in batch: token_ids = sentence_split(text) token_ids = sequence_padding(token_ids, length=maxlen) batch_token_ids.append(token_ids) batch_labels.append(label) batch_token_ids = torch.tensor(sequence_padding(batch_token_ids, length=max_segment), dtype=torch.long, device=device) batch_labels = torch.tensor(batch_labels, device=device) return batch_token_ids, batch_labels class Attention(nn.Module): """注意力层。""" def __init__(self, hidden_size, **kwargs): self.hidden_size = hidden_size super().__init__(**kwargs) self.weight = nn.Linear(self.hidden_size, self.hidden_size, bias=False) self.bias = nn.Parameter(torch.zeros(self.hidden_size)) self.query = nn.Linear(self.hidden_size, 1, bias=False) def forward(self, x, mask): '''x: [btz, max_segment, hdsz] mask: [btz, max_segment, 1] ''' mask = mask.squeeze(2) # [btz, max_segment] # linear key = self.weight(x) + self.bias # [btz, max_segment, hdsz] # compute attention outputs = self.query(key).squeeze(2) # [btz, max_segment] outputs -= 1e32 * (1 - mask) attn_scores = F.softmax(outputs, dim=-1) attn_scores = attn_scores * mask attn_scores = attn_scores.reshape(-1, 1, attn_scores.shape[-1]) # [btz, 1, max_segment] outputs = torch.matmul(attn_scores, key).squeeze(1) # [btz, hdsz] return outputs # 定义bert上的模型结构 class Model(BaseModel): def __init__(self): super().__init__() self.bert = build_transformer_model(config_path=config_path, checkpoint_path=checkpoint_path, segment_vocab_size=0) self.dropout1 = nn.Dropout(0.1) self.dropout2 = nn.Dropout(0.1) self.attn = Attention(768) self.dense = nn.Linear(768, num_classes) def forward(self, token_ids): ''' token_ids: [btz, max_segment, max_len] ''' input_mask = torch.any(token_ids, dim=-1, keepdim=True).long() # [btz, max_segment, 1] token_ids = token_ids.reshape(-1, token_ids.shape[-1]) # [btz*max_segment, max_len] output = self.bert([token_ids])[:, 0] # [btz*max_segment, hdsz] output = output.reshape((-1, max_segment, output.shape[-1])) # [btz, max_segment, hdsz] output = output * input_mask output = self.dropout1(output) output = self.attn(output, input_mask) output = self.dropout2(output) output = self.dense(output) return output class Evaluator(Callback): def __init__(self, model, dataloader, fold): super().__init__() self.model = model self.dataloader = dataloader self.best_val_f1 = 0. self.fold = fold def evaluate(self): y_true, y_pred = list(), list() for x, y in tqdm(self.dataloader, desc='evaluate'): y_true.append(y.cpu().numpy()) y_pred.append(self.model.predict(x).argmax(axis=1).cpu().numpy()) y_true = np.concatenate(y_true) y_pred = np.concatenate(y_pred) f1 = f1_score(y_true, y_pred, average='macro') return f1 def on_epoch_end(self, steps, epoch, logs=None): val_f1 = self.evaluate() if val_f1 > self.best_val_f1: self.best_val_f1 = val_f1 self.model.save_weights(f'best_model_fold{self.fold}.pt') logs['val_f1'] = val_f1 # 这个要设置，否则EarlyStopping不生效 print(f'val_f1: {val_f1:.5f}, best_val_f1: {self.best_val_f1:.5f}\n') def do_train(df_train): skf = StratifiedKFold(n_splits=n, random_state=SEED, shuffle=True) for fold, (trn_idx, val_idx) in enumerate(skf.split(df_train['text'], df_train['label']), 1): print(f'[Fold {fold}]') train_data = load_data(df_train.iloc[trn_idx]) valid_data = load_data(df_train.iloc[val_idx]) train_dataloader = DataLoader(ListDataset(data=train_data), batch_size=batch_size, shuffle=True, collate_fn=collate_fn) valid_dataloader = DataLoader(ListDataset(data=valid_data), batch_size=batch_size, collate_fn=collate_fn) model = Model().to(device) model.compile(loss=nn.CrossEntropyLoss(), optimizer=optim.Adam(model.parameters(), lr=lr), adversarial_train={'name': 'fgm'}) callbacks = [ Evaluator(model, valid_dataloader, fold), EarlyStopping(monitor='val_f1', patience=5, verbose=1, mode='max'), # 需要在Evaluator后面 ] model.fit( train_dataloader, steps_per_epoch=None, epochs=epochs, grad_accumulation_steps=grad_accum_steps, callbacks=callbacks ) del model if __name__ == '__main__': df_train = pd.read_csv('E:/Github/天池新闻分类/data/train_set.csv', sep='\t') df_train['text'] = df_train['text'].apply(lambda x: x.strip().split()) do_train(df_train)