Python - 深度学习系列38 重塑实体识别5-预测并行化改造

说明

在重塑实体识别4中梳理了数据流，然后我发现pipeline的串行效率太低了，所以做了并行化改造。里面还是有不少坑的，记录一下。

内容

1 pipeline

官方的pipeline看起来的确是比较好用的，主要是实现了比较好的数据预处理。因为在训练/使用过程中都要进行数据的令牌化与反令牌化，有些字符会被特殊处理，例如 '##A’等。

在使用过程中，我用200条新闻数据进行测试，用pipeline方法花了11分钟处理完毕，期间CUDA的使用率大约为10%。按此估算，即使用多接口并行的方式，那么一分钟最多处理2000条，一天最多处理0.14*2000~30万条数据。这个效率太低了。

2 并行化

最终的结论是不到30秒处理200条，显存只占用2.6G，理论上可以支持3个服务并行(以确保GPU的完全利用)。按最保守的估计，改造后的并行化应该可以提升3倍的效率，稍微激进一点，可以提升10倍的效率。这个之后可以进行测试。

一些主要的点如下

2.1 结果解析

结果可以分为：

• 1 仅含解析出的实体列表，用逗号连接字符串表示。
• 2 含实体及其起始位置的表示，这个用于标注反馈、二次增强处理。
• 3 仅含BIO标签，主要用于和测试数据进行效果比对。

对应的相关函数，看起来有点繁杂，我自己都不太想看第二眼。

from datasets import ClassLabel
# 定义标签列表
label_list = ['B', 'I', 'O']
# 创建 ClassLabel 对象
class_label = ClassLabel(names=label_list)
def convert_entity_label_batch(x):
    x1 = x
    return class_label.int2str(x1)
# 定义函数将整数Tensor转换为字符串 | 反令牌函数，但是用不上；因为predict label列表的长度和 ss_padding相同
def tensor_to_string(tensor, tokenizer = None , skip_special_tokens = True):
    return tokenizer.decode(tensor.tolist(), skip_special_tokens=skip_special_tokens).replace(' ','')

from datasets import ClassLabel
def detokenize(word_piece):
    """
    将 WordPiece 令牌还原为原始句子。
    """
    if word_piece.startswith('##'):
        x = word_piece[2:]
    else:
        x = word_piece
    return x
import re
def extract_bio_positions(bio_string):
    pattern = re.compile(r'B(I+)(O|$)')
    matches = pattern.finditer(bio_string)
    
    results = []
    for match in matches:
        start, end = match.span()
        results.append((start, end - 1))  # end-1 to include the last 'I'
    
    return results

# 0.1ms
def parse_ent_pos_map_batch(some_dict = None):

    word_list = some_dict['token_words']
    label_list = [int(x) for x in list(some_dict['label_list'])]

    min_len = min(len(word_list),len(label_list))
    word_list = word_list[:min_len]
    label_list = label_list[:min_len]
    

    label_list1 =  list(map(convert_entity_label_batch,label_list))
    oriword_list1 = list(map(detokenize,word_list))
    ori_word_str =''.join(oriword_list1)
    # 补到等长
    label_str = ''
    for i in range(len(label_list1)):
        len_of_ori_word = len(oriword_list1[i])
        if len_of_ori_word == 1:
            tem_str = label_list1[i]
        else:
            if label_list1[i] in ['I','O']:
                tem_str = label_list1[i] * len_of_ori_word
            else:
                tem_str = 'B' + 'I' * (len_of_ori_word -1)        
        label_str += tem_str

    pos_list = extract_bio_positions(label_str)
    part_ent_list = [(ori_word_str[x[0]:x[1]] , *x) for x in pos_list]
    return part_ent_list

# =============
def make_BIO_by_len(some_len):
    default_str = 'I' * some_len
    str_list = list(default_str)
    str_list[0] ='B'
    return str_list
def gen_BIO_list2(some_dict):
    the_content = some_dict['clean_data']
    ent_list =  some_dict['ent_tuple_list']

    content_list = list(the_content)
    tag_list = list('O'* len(content_list))
    
    for ent_info in ent_list:
        start = ent_info[1]
        end = ent_info[2]
        label_len = end-start
        tem_bio_list = make_BIO_by_len(label_len)
        tag_list[start:end] = tem_bio_list

    res_dict = {}
    res_dict['x'] = content_list
    res_dict['y'] = tag_list
    return res_dict

def trim_len(some_dict = None):
    padding_BIO = some_dict['padding_BIO']
    ss_len = some_dict['ss_len']
    return padding_BIO[:ss_len]

2.2 批量预测

看起来同样很繁杂，但是不得不细看。首先，数据会按照几个长度 20，50，198分为三部分处理，batch_predict每次仅处理一个批次。在这里，将数据转为定长的令牌长度，然后转入CUDA进行批量预测。结果再按照实体-位置 tuple， 实体列表和BIO三种方式进行解析。

from functools import partial
import transformers 
import torch 
from transformers import AutoModelForMaskedLM, AutoTokenizer,AutoModelForTokenClassification
from functools import partial
# some_batch 是原文经过padding的数据，['ss_hash','ss','ss_len', 'ss_padding'], 其中ss_padding的长度是固定的
# 模型文件和令牌文件都放在model_path之下，model比较大，避免重载；而tokenize会有padding过程，必须重载
# 模型先载入cuda
def batch_predict(some_batch, ss_padding_len = None, model = None, model_path = None):
    # 因为tokenize会在令牌的前后加上分隔令牌，所以+2
    if ss_padding_len is None:
        ss_padding_len = some_batch['ss_padding'].apply(len).max()
    print('ss_padding_len is %s ' % ss_padding_len)
    max_len = ss_padding_len+2
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tencoder = partial(tokenizer.encode,truncation=True, max_length=max_len, is_split_into_words=True, return_tensors="pt",  padding='max_length')
    some_batch['ss_padding_token'] = some_batch['ss_padding'].apply(list).apply(tencoder)

    # 构成矩阵
    minput = torch.cat(list(some_batch['ss_padding_token'].values))
    # 将数据搬到GPU中处理再返回
    with torch.no_grad():
        input_cuda = minput.to(device)
    
        outputs_cuda = model(input_cuda).logits
        predictions = torch.argmax(outputs_cuda, dim=2)
    
        predictions_list = list(predictions.to('cpu').numpy())
    
    predict_list1 = []
    for predictions in predictions_list:
        tem_pred_tag = [int(x) for x in predictions[1:-1]]
        predict_list1.append(tem_pred_tag)

    some_batch['label_list'] = predict_list1
    _s = cols2s(some_df =some_batch, cols= ['ss_padding','label_list'], cols_key_mapping= ['token_words', 'label_list'])
    _s1 = _s.apply(parse_ent_pos_map_batch)
    some_batch['ent_tuple_list'] = list(_s1)
    some_batch['ent_list'] = some_batch['ent_tuple_list'].apply(lambda x: ','.join([a[0] for a in x ]))

    _s = cols2s(some_batch, cols= ['ss_padding', 'ent_tuple_list'], cols_key_mapping= ['clean_data', 'ent_tuple_list'])
    s1 = _s.apply(gen_BIO_list2)
    ent_tuple_res_df1 = pd.DataFrame(s1.to_list())
    some_batch['padding_BIO'] = list(ent_tuple_res_df1['y'].apply(lambda x: ''.join(x)))
    _s00 = cols2s(some_batch, cols = ['ss_len', 'padding_BIO'], cols_key_mapping=['ss_len', 'padding_BIO'])
    some_batch['BIO'] = list(_s00.apply(trim_len))
    return some_batch    

3 迭代器

在推送数据处理时，可以采用迭代器来控制不同的批次数据

# 迭代器切分
import pandas as pd
class DataFrameBatchIterator:
    def __init__(self, dataframe, batch_size):
        self.dataframe = dataframe
        self.batch_size = batch_size
        # 【我增加的】
        self.fail_batch_list = []
    def __iter__(self):
        num_rows = len(self.dataframe)
        num_batches = (num_rows - 1) // self.batch_size + 1
        for i in range(num_batches):
            start_idx = i * self.batch_size
            end_idx = (i + 1) * self.batch_size
            batch_data = self.dataframe.iloc[start_idx:end_idx]
            yield batch_data
    # 【我增加的】
    def clear_fail(self):
        self.fail_batch_list = []
    # 【我增加的】
    def get_some_batch(self, batch_idx):
        return self.dataframe.iloc[self.batch_size * batch_idx: self.batch_size * (batch_idx + 1)]
    # 【我增加的】记录失败的批次
    def rec_fail_batch_idx(self, batch_idx):
        self.fail_batch_list.append(batch_idx)
# 创建一个示例 DataFrame
data = {'Name': ['John', 'Jane', 'Mike', 'Alice', 'Bob'],
        'Age': [25, 30, 35, 28, 32],
        'City': ['New York', 'Paris', 'London', 'Tokyo', 'Sydney']}
df = pd.DataFrame(data)
# 创建 DataFrame 迭代器
batch_iterator = DataFrameBatchIterator(df, batch_size=2)
import tqdm
# 使用迭代器逐批次处理数据
for i,batch in tqdm.tqdm(enumerate(batch_iterator)):
    try:
        # 在这里可以对当前批次的数据进行相应的操作
        # 例如进行数据清洗、特征处理、模型训练等
        # 示例：打印当前批次的数据
#         raise Exception(e) 
        print(batch)
    except:
        print('>>> %s Fail' % i)
        batch_iterator.rec_fail_batch_idx(i)

以下是实际的调度

# 假设处理长度为1万的句子
# 20 * 2000 ~ 4w
# 50 * 800 ~  4w
# 200 * 200 ~ 4w
import warnings 
warnings.filterwarnings('ignore')
batch_slice_para = {20:2000, 50:800, 200:200}
batch_len_list = sorted(list(batch_slice_para.keys()))
batch_len_list.insert(0,0)

batch_df_list = []
for i in range(len(batch_len_list)):
    if i >0:
        sel = (ss_df['ss_len'] >=batch_len_list[i-1]) & (ss_df['ss_len'] < batch_len_list[i])
        if sel.sum():
            padding_len = batch_len_list[i]
            padding_batch = batch_slice_para[padding_len]
            
            tem_df= ss_df[sel]
            # tem_df['ss_padding'] = tem_df['ss'].apply(lambda x: x.ljust(padding_len,'a'))
            tem_df['ss_padding'] = tem_df['ss']
            tem_df_iterator = DataFrameBatchIterator(tem_df, padding_batch)
            batch_df_list.append(tem_df_iterator)
        else:
            batch_df_list.append(None)

对每个批次执行处理，载入模型

label_list = ['B','I','O']
model_checkpoint = 'model03'
model = AutoModelForTokenClassification.from_pretrained(model_checkpoint, num_labels=len(label_list))

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Device: %s' % device)
# 自动切换设备
if model.device.type != device:
    model.to(device)
    print('>>> 检测到模型设备与当前指定不一致，切换 %s' % device )
else:
    print('>>> 模型设备一致，不切换 %s' % device)

分批次预测(主要是确保显存不溢出)

batch_res_list = []
for some_iter in batch_df_list:
    for some_batch in some_iter:
        batch_res = batch_predict(some_batch, model = model, model_path = 'model03')
        batch_res_list.append(batch_res)

结果合并

batch_res_df = pd.concat(batch_res_list, ignore_index= True)
mdf = pd.merge(input_df , batch_res_df[['ss_hash', 'ent_list']],how='left', on ='ss_hash')

4 总结

一个在理论上证明可以显著提升效率的点在于，模型进行实体识别时先切分了短句，然后按短句进行了去重：相同短句的实体结果一定是相同的。

实验中，200条新闻产生了约5万个短句，去重后只剩下约3.5万。所以即使在这一步也是有提升的。当然，这种方式同样也可以被用于pipeline。

还有就是在处理填充时，并不按照最大长度统一填充。而是按照句子长度的统计特性分为了短、中、长三种方式。从统计上看，约70%的短句长度是在20个字符以内的，真正超过50个字符的短句(中间无分隔符)，即使从语法上来看也是比较奇怪的。
这样在填充数据时浪费就比pipeline要小，同样显存可以装下更多的数据。