python使用多进程加快处理数据

最近需要用python循环遍历xlsx文件，但总感觉特慢，于是记录一下使用多进程让cpu性能拉满，加快速度。

1.原始代码

下面展示 原始代码。

这个举例代码的是用来循环遍历我文件夹下的每一个xlsx文件，
并读取时间以及对应的电流信息，
并将label为0以及为1的时间和对应的电流分别读取保存。

'''这个函数文件是将把时间加进去，将时间和电流转换为二维矩阵，构建：时域特征'''
import os

import pandas as pd
import numpy as np

file_dir = r'D:\交流并联（正常=0有弧=1）'
file_dir_list = os.listdir(file_dir)
final_data_0 = []  # 初始化一个空列表用于存储每次循环得到的data_0
final_data_1 = []  #
# 循环遍历xlsx文件
for xlsx in file_dir_list:
    curr_xlsx = os.path.join(file_dir, xlsx)
    # 读取xlsx文件
    df = pd.read_excel(curr_xlsx, engine='openpyxl')
    # 读取xlsx文件列数据
    time_data = np.array(df['TIME'])
    current_data = np.array(df['Current'])
    label_data = np.array(df['Label'])
    #将label中为0对应的时间-电流（时域）数据筛选出来
    time_0 = time_data[label_data == 0]
    current_0 = current_data[label_data == 0]
    #将label中为1对应的时间-电流（时域）数据筛选出来
    time_1 = time_data[label_data == 1]
    current_1 = current_data[label_data == 1]

    data_0 = np.concatenate((np.expand_dims(time_0, 1), np.expand_dims(current_0, 1)), axis=1)
    data_1 = np.concatenate((np.expand_dims(time_1, 1), np.expand_dims(current_1, 1)), axis=1)
    final_data_0.append(data_0)  # 将当前循环得到的data_0添加到final_data_0列表
    final_data_1.append(data_1)
    print(data_0.shape)
    print("H_W_0 is :{}".format(np.sqrt(data_0.shape[0])))
    print(data_1.shape)
    print("H_W_1 is :{}".format(np.sqrt(data_1.shape[0])))

    print('-----------')
# 将所有循环得到的data_0和data_1按照第一个维度进行拼接
final_data_0 = np.concatenate(final_data_0, axis=0)
final_data_1 = np.concatenate(final_data_1, axis=0)

2.使用多进程改写代码

下面展示 改写使用多进程的代码。

其实主要还是这个代码：
# 创建进程池
    with Pool(processes=os.cpu_count()) as pool:
        results = pool.map(process_file, file_dir_list)
流程也就是：
（1）将原始代码封装成函数（不要用循环）
（2）pool.map(process_file, file_dir_list)第一参函数就是（1）封装好的函数，
然后第二个参数就是原始代码中循环的每个xlsx文件名。
记住这个流程，你就学会多进程的使用啦！
     

import os
import pandas as pd
import numpy as np
from multiprocessing import Pool

# 定义处理每个文件的函数
def process_file(xlsx):
    file_dir = r'D:\交流并联（正常=0有弧=1）'
    curr_xlsx = os.path.join(file_dir, xlsx)
    df = pd.read_excel(curr_xlsx, engine='openpyxl')
    time_data = np.array(df['TIME'])
    current_data = np.array(df['Current'])
    label_data = np.array(df['Label'])

    time_0 = time_data[label_data == 0]
    current_0 = current_data[label_data == 0]
    data_0 = np.concatenate((np.expand_dims(time_0, 1), np.expand_dims(current_0, 1)), axis=1)

    time_1 = time_data[label_data == 1]
    current_1 = current_data[label_data == 1]
    data_1 = np.concatenate((np.expand_dims(time_1, 1), np.expand_dims(current_1, 1)), axis=1)

    return (data_0, data_1)


# 使用多进程读取和处理文件
def main():
    file_dir = r'D:\交流并联（正常=0有弧=1）'
    file_dir_list = os.listdir(file_dir)

    # 创建进程池
    with Pool(processes=os.cpu_count()) as pool:
        results = pool.map(process_file, file_dir_list)

    final_data_0 = np.concatenate([result[0] for result in results], axis=0)
    final_data_1 = np.concatenate([result[1] for result in results], axis=0)

    print('final_data_0 shape:', final_data_0.shape)
    print('final_data_1 shape:', final_data_1.shape)


if __name__ == '__main__':
    main()