Python多进程(multiprocessing)
最近在做一些联邦学习方面的研究,因为联邦学习的过程涉及到多个客户端训练,一开始在写程序的时候自己并没有在程序中加入多进程相关的代码,导致训练神经网络的速度比较慢,服务器一直是“一核有难,多核围观”😂后来查阅了一些相关的资料,在程序中加入了多进程,速度提升了不少,在这里记录一下一些自己调试的经验。
首先是为什么要使用多进程而不是多线程。因为Python的GIL(全局解释器锁)的存在,导致Python的多线程是虚假的多线程,所以如果要进行并行运算,我们要使用多进程来达到更好的效果
在这次的Python多进程实践中,我主要使用的是multiprocessing
这个类。我们先来看一个简单的例子
Python多进程训练神经网络:
from datetime import time
import torch.multiprocessing as mp
import torch
from model import MNIST_Net
from torchvision import datasets, transforms
def train(model, data_loader):
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for data, labels in data_loader:
data = data.to(device)
labels = labels.to(device)
optimizer.zero_grad()
loss = torch.nn.functional.cross_entropy(model(data), labels)
loss = loss.to(device)
loss.backward()
optimizer.step()
print(loss)
if __name__ == '__main__':
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device_cpu = "cpu"
train_dataset = datasets.MNIST("", train=True, download=True, transform=transforms.ToTensor())
eval_dataset = datasets.MNIST("", train=False, transform=transforms.ToTensor())
data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=512)
model_1 = MNIST_Net()
model_2 = MNIST_Net()
model_1 = model_1.to(device)
model_2 = model_2.to(device)
model = [model_1, model_2]
optimizer_model_1 = torch.optim.Adam(model_1.parameters(), lr=0.001)
optimizer_model_2 = torch.optim.Adam(model_2.parameters(), lr=0.001)
processes = []
process = [mp.Process(target=train, args=(m, data_loader)) for m in model]
[p.start() for p in process]
[p.join() for p in process]
假设我们有model_1和model_2两个神经网络,我们希望能够将两个神经网络放到不同的进程里面进行训练,来提高训练的速度。整段代码的核心在最后四行,我们用代码分别为这两个模型的训练创建了一个进程,从而达到了并行化的效果。
需要注意的是如果你使用的是Windows而不是Linux系统的话,如果用了cuda就会报错(如上图所示),解决办法是把程序放到Linux上面跑。
如果我们要使用进程间通信,我们可以选择Queue,Pipe或者Pool,一般来说Pipe的速度最快,Queue次之,Pool的速度最慢。下面用一个简单的例子来展示一下其中Queue的用法。
import multiprocessing as mp
from multiprocessing import Queue, Process
def test(queue_1):
queue_1.put(1)
if __name__ == '__main__':
candidates = [1, 2, 3]
queue_1 = Queue()
process = [Process(target=test, args=(queue_1,)) for c in candidates]
[p.start() for p in process]
[p.join() for p in process]
res = [queue_1.get() for p in process]
print(res)
结果:
从结果我们可以看出程序的多进程的运行正确。
Python多进程中遇到bug,一些可能的解决方案:
- 去掉多进程中的方法的return
- 使用
process = [Process(target=test(queue_1)) for c in candidates]
代替原本的那行代码(以上面的例子为例)