如何用Multiprocessing并行实现ResNet50图像向量化？

原多进程代码的核心问题 #

1. TensorFlow模型不能跨进程直接共享：你直接在主进程加载模型后让子进程调用，TensorFlow的模型对象没法在多进程间直接传递，会触发序列化错误或者运行异常。
2. 子进程内拼接DataFrame无效：每个子进程有独立的内存空间，你在子进程里修改的embedding_df根本同步不到主进程，而且每次返回拼接后的小df纯粹浪费内存。
3. 进程池用法错误：pool.map返回的是所有子任务结果的列表，你直接赋值给new_df，最后得到的是一堆小DataFrame的集合，不是合并好的大表。

修正后的多进程代码 #

下面是可正常运行的代码，关键改进都做了标注：

import os
import numpy as np
import pandas as pd
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from tensorflow.keras.preprocessing import image
from multiprocessing import Pool
from tqdm import tqdm

# 每个子进程启动时执行：独立加载模型
def init_worker():
    global model
    # 子进程单独加载模型，避开跨进程共享的坑
    model = ResNet50(include_top=False, weights='imagenet', pooling='avg')

# 单张图片处理逻辑，子进程会调用此函数
def process_image(img_name):
    img_path = os.path.join('food-101/images', img_name)
    try:
        img = image.load_img(img_path, target_size=(224, 224))
        x = image.img_to_array(img)
        x = np.expand_dims(x, axis=0)
        x = preprocess_input(x)
        preds = model.predict(x, verbose=0)  # 关闭预测日志，避免输出乱序
        # 返回元组（embedding数组，图片名），比返回DataFrame更省内存
        return (preds[0], img_name)
    except Exception as e:
        print(f"处理{img_name}出错: {e}")
        return (None, img_name)

if __name__ == '__main__':
    # 获取所有图片文件名
    images_path = os.listdir('food-101/images/')
    
    # 创建进程池，指定初始化函数，每个子进程启动时自动加载模型
    with Pool(processes=os.cpu_count(), initializer=init_worker) as pool:
        # 用tqdm显示实时处理进度
        results = list(tqdm(pool.imap(process_image, images_path), total=len(images_path)))
    
    # 过滤处理失败的图片结果
    valid_results = [res for res in results if res[0] is not None]
    
    # 将结果转换为DataFrame
    embeddings = np.array([res[0] for res in valid_results])
    img_names = [res[1] for res in valid_results]
    
    embedding_df = pd.DataFrame(embeddings)
    embedding_df['image'] = img_names
    
    # 保存到CSV文件
    embedding_df.to_csv('./embeddings.csv', index=False, header=True)

关键改进说明 #

• 子进程独立加载模型：通过initializer=init_worker让每个进程池内的子进程自己加载模型，这是TensorFlow多进程的标准操作，彻底解决模型跨进程共享的问题。
• 高效返回结果：子进程返回元组而非DataFrame，减少内存消耗和序列化时间，最后统一合并成大DataFrame，比原代码循环中反复pd.concat效率高很多。
• 异常捕获机制：单张图片损坏或加载失败时，程序不会直接崩溃，还能打印出错的图片名方便排查问题。
• 进度可视化：用pool.imap配合tqdm，能实时看到处理进度，比map更直观。

额外提速建议 #

1. 批量处理图片：如果内存充足，可以修改代码让每个子进程处理一批图片，减少model.predict的调用次数——批量预测比单张预测效率高很多，能进一步压缩处理时间。
2. 尝试TensorFlow Dataset：要是不想用多进程，也可以用TensorFlow原生的tf.data.Dataset构建图片数据集，配合模型批量预测，这也是TensorFlow官方推荐的高效并行方式。

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