python 多进程执行任务

python 多进程执行任务

最近在写分布式优化, 想要实现简单的并行程序, 实现Parameter Server结构的一类sgd算法
看他们收敛性. 并行可以充分利用CPU/GPU资源, 加快训练. 对于一些简单的任务却有多个baseline, 我们就可以并行跑, 调试贼爽.
这里主要用到python中的Pipe来进行通讯

代码实现

from multiprocessing import *
from numpy import *

def f(conn):
    while True:
        res = conn.recv()
        if res is None:
            conn.close()
            break
        conn.send(res*res)


if __name__ == "__main__":
    parent_conn_list = []
    child_conn_list = []
    work_list = []
    for i in range(5):
        parent_conn,child_conn = Pipe()
        parent_conn_list.append(parent_conn)
        child_conn_list.append(child_conn)  
        p = Process(target = f, args = (child_conn,))
        p.start()
        work_list.append(p)

    ans = 0
    print("The number of CPU is:" + str(cpu_count()))
    for i in range(10):
        for j in range(5):
            parent_conn = parent_conn_list[j]
            parent_conn.send(i*5+j)

        for j in range(5):
            parent_conn = parent_conn_list[j]
            ans += parent_conn.recv()
    print("ans ", ans)
    ans0 = 0
    for i in range(50):
        ans0 += i*i
    print("ans0 ",ans0)
    for conn in parent_conn_list:
        conn.send(None)
        conn.close()
    for work in work_list:
        work.join()

结果如下

The number of CPU is:24
ans 40425
ans0 40425