Python使用多进程生成BERT词嵌入失败问题求助

单进程正常代码 #

以下单进程代码可正常生成BERT词嵌入：

preprocess = 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3'
small_bert = 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/1'

bert_pre = hub.KerasLayer(preprocess)
bert_model = hub.KerasLayer(small_bert)

def word_embedding(sent):
    if type(sent)==str:
        sent = [sent]
    return bert_model(bert_pre(sent))['pooled_output']

调用示例输出正常：

word_embedding("This is a dog")

<tf.Tensor: shape=(1, 128), dtype=float32, numpy=
array([[-0.9999982 ,  0.09135217, -0.9993329 ,  0.96194535, -0.9991869 ,
         0.07579318, -0.9892281 , -0.9668794 ,  0.06745824,  0.07896714,
        -0.81101316,  0.01557093, -0.11514608,  1.        , -0.9598647 ,
        -0.83923537,  0.8639652 ,  0.04229677, -0.9274309 ,  0.87699145,
         0.9675325 ,  0.02446483,  0.96745443,  0.9137065 , -0.99994516,
        -0.00484818, -0.9999251 ,  0.97071564,  0.9577635 ,  0.12064736,
         0.13986832,  0.01977904, -0.99208915,  0.1142559 ,  0.98749965,
         0.9999112 , -0.93131316, -0.05863096,  0.9166601 , -0.9995932 ,
         0.92014617,  0.94889516, -0.9995392 ,  0.9891672 , -0.9999985 ,
        -0.15958041, -0.99989873,  0.9984788 ,  0.9674587 ,  0.9849434 ,
         0.9884202 , -0.5348996 ,  0.07992988,  0.9977897 ,  0.99813277,
         0.9999679 , -0.99952024, -0.970687  ,  0.9040054 , -0.9458808 ,
         0.01751512,  0.36849502,  0.3939357 ,  0.9101836 , -0.15718344,
        -0.99999946, -0.44872832, -0.6077106 ,  0.96371555,  0.5564301 ,
         0.9982054 , -0.09481359, -0.9996993 ,  0.03875534,  0.65399134,
        -0.9902064 ,  0.66297245,  0.10515413, -0.97484773,  0.18679208,
        -0.5837009 , -0.12993163, -0.96478623, -0.99981767,  0.99985546,
        -0.98870945,  0.8561884 , -0.58723676, -0.68301636,  0.67417735,
        -0.9766185 ,  0.9956491 , -0.88204795,  0.99866074,  0.2829505 ,
         0.42085564, -0.9546872 , -0.8894943 , -0.9999068 , -0.97645766,
        -0.99447215,  0.97132486, -0.9995873 , -0.90443873, -0.9787839 ,
        -0.6670069 , -0.9991659 , -0.9913582 , -0.19619215,  0.9979996 ,
         0.99873877,  0.94075304, -0.76902175,  0.9997495 , -1.        ,
         0.06643485,  0.8816498 ,  0.83833504,  0.09686996, -0.9954674 ,
         0.22044522, -0.99998134, -0.5231443 ,  0.902108  , -0.9998227 ,
         0.97717   ,  0.9373147 ,  0.9990008 ]], dtype=float32)>

多进程问题代码 #

尝试用多进程并行处理时，进程能从队列获取元素，但卡在word_embedding步骤无响应，且无CPU负载：

from multiprocessing import Queue, Process
from multiprocessing.queues import Empty
pending=Queue()
for x in  list(model_pdf['JCOM']):
    pending.put(x)
def job():
    print("Started")
    kill=0
    while pending.qsize()>0:
        try:
            x = pending.get(block=True,timeout=2)
            print("Can I get some element?")
            word_embedding(x)
            print("Can I finish some embedding?")
        except Empty:
            return 1
job_list=[]
for x in range(60):
    job_list.append(Process(target=job))
for j in job_list:
    j.start()
for j in job_list:
    j.join()

环境：Ubuntu 22.04 LTS，64核CPU+64GB内存

核心原因 #

TensorFlow的计算图和TF Hub的KerasLayer无法在多进程间共享，主进程加载的模型不能直接被子进程调用，会导致进程死锁或无响应。

具体修复步骤 #

1. 每个子进程独立加载模型：将BERT预处理层和模型的加载逻辑放到job函数内部，确保每个进程拥有独立的模型实例。
2. 合理设置进程数：60个进程远超实际可用资源，建议设置为CPU核心数（如os.cpu_count()），避免上下文切换开销。
3. 优化队列判断逻辑：pending.qsize()在多进程中并非线程安全，改用get的阻塞特性循环处理，直到队列为空。

修正后的代码 #

from multiprocessing import Queue, Process
import os
import tensorflow_hub as hub

pending = Queue()
for x in list(model_pdf['JCOM']):
    pending.put(x)

def job():
    # 每个进程独立加载模型
    preprocess = 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3'
    small_bert = 'https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/1'
    bert_pre = hub.KerasLayer(preprocess)
    bert_model = hub.KerasLayer(small_bert)

    def word_embedding(sent):
        if type(sent) == str:
            sent = [sent]
        return bert_model(bert_pre(sent))['pooled_output']

    print(f"Process {os.getpid()} started")
    while True:
        try:
            x = pending.get(block=True, timeout=2)
            print(f"Process {os.getpid()} processing: {x}")
            embedding = word_embedding(x)
            # 这里可以添加保存或处理embedding的逻辑
            print(f"Process {os.getpid()} finished processing: {x}")
        except Empty:
            print(f"Process {os.getpid()} exiting")
            break

# 设置进程数为CPU核心数
process_count = os.cpu_count()
job_list = []
for _ in range(process_count):
    job_list.append(Process(target=job))

for j in job_list:
    j.start()

for j in job_list:
    j.join()

额外优化建议 #

• 模型缓存：TF Hub会自动缓存下载的模型，无需担心重复下载占用带宽。
• 批量处理：如果单条文本处理效率低，可以修改队列传递批量文本，减少进程间通信开销。
• 使用进程池：改用multiprocessing.Pool可以简化代码，自动管理进程生命周期。

内容的提问来源于stack exchange，提问作者Bayesian