Running multiple inferences in parallel with PyTorch

Question

I'm trying to implement Double DQN (not to be confused with DQN with a slightly delayed Q-target network) in PyTorch to train an agent to play an Atari OpenAI Gym game. Here I discuss the implementation of the following formula: Update of Q-network, formula taken from Sutton & Barto.

My first implementation is:

Q_pred = self.Q_1.forward(s_now)[T.arange(batch_size), actions.long()]

Q_next_all = self.Q_1.forward(s_next)
maxA_id = T.argmax(Q_next_all, dim=1)
Q_pred2 = self.Q_2.forward(s_next)[T.arange(batch_size), maxA_id]

Q_target = (rewards + (~dones) * self.GAMMA * Q_pred2).detach()

self.Q_1.optimizer.zero_grad()
self.Q_1.loss(Q_target, Q_pred).backward()
self.Q_1.optimizer.step()

(Q_1 and Q_2 are nn.Module classes, and all of the variables involved here are already torch tensors lying in the GPU.)

I noticed that my program ran much slower than a previous implementation which used plain DQN.

I realized that I can combine the batches entering Q_1, so there will be one combined batch being forwarded in the neural network, instead of two batches in sequence. The code becomes:

s_combined = T.cat((s_now, s_next))
Q_combined = self.Q_1.forward(s_combined)

Q_pred = Q_combined[T.arange(batch_size), actions.long()]
Q_next_all = Q_combined[batch_size:]

Q_pred2_all = self.Q_2.forward(s_next)

maxA_id = T.argmax(Q_next_all, dim=1)
Q_pred2 = Q_pred2_all[T.arange(batch_size), maxA_id]
Q_target = (rewards + (~dones) * self.GAMMA * Q_pred2).detach()

self.Q_1.optimizer.zero_grad()
self.Q_1.loss(Q_target, Q_pred).backward()
self.Q_1.optimizer.step()

(This proves that I understand how to do batch training in PyTorch, so don't mark this as a duplicate of this question.)

Furthermore, I realized that Q_1 and Q_2 can process their batches in parallel. So I looked up how to do multiprocessing in PyTorch. Unfortunately, I couldn't find a good example. I tried to adapt a code that looks similar to my scenario, and my code becomes:

def spawned():
    s_combined = T.cat((s_now, s_next))
    Q_combined = self.Q_1.forward(s_combined)

    Q_pred = Q_combined[T.arange(batch_size), actions.long()]
    Q_next_all = Q_combined[batch_size:]

mp.set_start_method('spawn', force=True)
p = mp.Process(target=spawned)
p.start()
Q_pred2_all = self.Q_2.forward(s_next)
p.join()

maxA_id = T.argmax(Q_next_all, dim=1)
Q_pred2 = Q_pred2_all[T.arange(batch_size), maxA_id]
Q_target = (rewards + (~dones) * self.GAMMA * Q_pred2).detach()

self.Q_1.optimizer.zero_grad()
self.Q_1.loss(Q_target, Q_pred).backward()
self.Q_1.optimizer.step()

This crashes with the error message:

AttributeError: Can't pickle local object 'Agent.learn..spawned'

So how do I make this work?

(Achieving this in CUDA programming is trivial. One simply launches two device kernels using a sequential host code, and the two kernels are automatically computed in parallel in the GPU.)

Running multiple inferences in parallel with PyTorch

0 Answers