使用pool.map()处理列表的列表时报错，如何正确实现并行？

我计划用sentence_transformers计算标题相似度，借助multiprocessing.Pool的map方法实现并行处理。传入的data是列表的列表，但执行时抛出TypeError: 'int' object is not subscriptable错误，请问如何正确传递列表的列表给pool.map()以实现并行处理？

原代码 #

import sentence_transformers
import multiprocessing
from tqdm import tqdm
from multiprocessing import Pool

embedding_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')

data = [[100227, 7382501.0, 'view', 30065006, False, ''],
 [100227, 7382501.0, 'view', 57072062, True, ''],
 [100227, 7382501.0, 'view', 66405922, True, ''],
 [100227, 7382501.0, 'view', 5221475, False, ''],
 [100227, 7382501.0, 'view', 63283995, True, '']]

# 定义并行执行的函数
def process_data(chunk):
    results = []
    for row in chunk:
        print(row[0])
        work_id = row[1]
        mentioning_work_id = row[3]
        print(work_id)

        if work_id in df_text and mentioning_work_id in df_text:
            title1 = df_text[work_id]['title']
            title2 = df_text[mentioning_work_id]['title']
            embeddings_title1 = embedding_model.encode(title1,convert_to_numpy=True)
            embeddings_title2 = embedding_model.encode(title2,convert_to_numpy=True)
            
            similarity = np.matmul(embeddings_title1, embeddings_title2.T)
            
            results.append([row[0],row[1],row[2],row[3],row[4],similarity])
        else:
            continue
    return results

# 定义使用的CPU核心数
num_cores = multiprocessing.cpu_count()

# 分割数据为多个块
chunk_size = len(data) // num_cores
# chunks = [data[i:i+chunk_size] for i in range(0, len(data), chunk_size)]

# 创建进程池
pool = multiprocessing.Pool(processes=num_cores)

results = []
with tqdm(total=len(data)) as pbar:
    for i, result_chunk in enumerate(pool.map(process_data, data)):
        # 更新进度条
        pbar.update()
        # 将结果加入列表
        results += result_chunk

# 合并结果
final_result = results

报错信息 #

---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/multiprocessing/pool.py", line 121, in worker
    result = (True, func(*args, **kwds))
  File "/opt/conda/lib/python3.7/multiprocessing/pool.py", line 44, in mapstar
    return list(map(*args))
  File "<ipython-input-4-3aab73406a3b>", line 18, in process_data
    print(row[0])
TypeError: 'int' object is not subscriptable
"""

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
<ipython-input-4-3aab73406a3b> in <module>
     46 results = []
     47 with tqdm(total=len(data)) as pbar:
---> 48     for i, result_chunk in enumerate(pool.map(process_data, data)):
     49         # Update the progress bar
     50         pbar.update()

/opt/conda/lib/python3.7/multiprocessing/pool.py in map(self, func, iterable, chunksize)
    266         in a list that is returned.
    267         '''
---> 268         return self._map_async(func, iterable, mapstar, chunksize).get()
    269 
    270     def starmap(self, func, iterable, chunksize=None):

/opt/conda/lib/python3.7/multiprocessing/pool.py in get(self, timeout)
    655             return self._value
    656         else:
---> 657             raise self._value
    658 
    659     def _set(self, i, obj):

TypeError: 'int' object is not subscriptable

错误原因 #

核心问题是直接把data传给了pool.map()：

• pool.map()会遍历传入的可迭代对象（这里是data），把每个元素（也就是data里的单个子列表，比如[100227, 7382501.0, ...]）作为参数传给process_data函数
• 但你的process_data函数是按“接收一个包含多个子列表的chunk，遍历chunk里的row”编写的，现在参数变成了单个row，执行for row in chunk时，会遍历这个子列表的每个元素（比如100227这个int），后续执行row[0]自然会报错，因为int类型不支持下标访问。

修复方案 #

1. 恢复数据分块代码 #

取消注释chunks的生成代码，把data分割成多个包含若干子列表的chunk：

chunks = [data[i:i+chunk_size] for i in range(0, len(data), chunk_size)]

2. 传入chunks而非data给pool.map() #

让process_data真正接收chunk（多个子列表的集合），这样遍历出来的row是正确的子列表。

3. 修复进度条更新逻辑 #

每个chunk处理的行数可能不同，要按当前chunk返回的结果长度更新进度条，而非每次默认加1。

4. 子进程内初始化模型（关键） #

sentence_transformers的模型对象无法跨进程序列化，所以要在子进程的process_data函数内初始化模型，避免序列化错误。

5. 添加if __name__ == '__main__': #

这是Windows系统下多进程的必要写法，避免重复初始化进程。

修复后的完整代码 #

import sentence_transformers
import multiprocessing
from tqdm import tqdm
import numpy as np

# 替换为你的实际df_text数据
df_text = {}

def process_data(chunk):
    # 子进程内初始化模型，避免跨进程序列化问题
    embedding_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
    results = []
    for row in chunk:
        work_id = row[1]
        mentioning_work_id = row[3]

        if work_id in df_text and mentioning_work_id in df_text:
            title1 = df_text[work_id]['title']
            title2 = df_text[mentioning_work_id]['title']
            embeddings_title1 = embedding_model.encode(title1, convert_to_numpy=True)
            embeddings_title2 = embedding_model.encode(title2, convert_to_numpy=True)
            
            similarity = np.matmul(embeddings_title1, embeddings_title2.T)
            
            results.append([row[0], row[1], row[2], row[3], row[4], similarity])
        else:
            continue
    return results

if __name__ == '__main__':
    data = [[100227, 7382501.0, 'view', 30065006, False, ''],
            [100227, 7382501.0, 'view', 57072062, True, ''],
            [100227, 7382501.0, 'view', 66405922, True, ''],
            [100227, 7382501.0, 'view', 5221475, False, ''],
            [100227, 7382501.0, 'view', 63283995, True, '']]

    num_cores = multiprocessing.cpu_count()
    # 确保chunk_size至少为1，避免数据量小于核心数时分块为空
    chunk_size = max(1, len(data) // num_cores)
    chunks = [data[i:i+chunk_size] for i in range(0, len(data), chunk_size)]

    with multiprocessing.Pool(processes=num_cores) as pool:
        results = []
        with tqdm(total=len(data)) as pbar:
            for result_chunk in pool.map(process_data, chunks):
                # 按当前chunk处理的有效行数更新进度条
                pbar.update(len(result_chunk))
                results += result_chunk

    final_result = results

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