I'm trying to implement something with pytorch. I have 2 GPUs and I want to train 2 models as below:
model0 = Mymodel().to('cuda:0')
model1 = Mymodel().to('cuda:1')
opt0 = torch.optim.Adam(model0.parameters(), lr=0.01)
opt1 = torch.optim.Adam(model0.parameters(), lr=0.01)
# 1.Forward data into model0 on GPU0
out = model0(x.to('cuda:0'))
# 2.Calculate the loss on model0, update model0's parameters
model0.loss.backward()
opt0.step()
opt0.zero_grad()
# 3.Use model0's output as input of model1 on GPU1
out = model1(out.to('cuda:1'))
# 4.Calculate the loss on model1, update model1's parameters
model1.loss.backward()
opt1.step()
opt1.zero_grad()
I want to train them simultaneously to speed up the whole procedure, but I think the code now will wait step 2(or 4) finished and finally do step 3(or 1). How can I implement my idea? Or which technique is I need(e.g. model parel, thread, multiprocessing...)?
I've consider some article like this, but I think there is some worng with the result, and I think it actually doesn't train models simultaneously.