title: Running multiple inferences in parallel with PyTorch
body: <p>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:
  <a href="https://i.stack.imgur.com/FjmrS.png" rel="nofollow noreferrer">Update of Q-network, formula taken from Sutton &amp; Barto</a>.</p>
  <p>My first implementation is:</p>
  <pre><code>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()
  </code></pre>
  <p>(Q_1 and Q_2 are nn.Module classes, and all of the variables involved here are already torch tensors lying in the GPU.)</p>
  <p>I noticed that my program ran much slower than a previous implementation which used plain DQN.</p>
  <p>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:</p>
  <pre><code>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()
  </code></pre>
  <p>(This proves that I understand how to do batch training in PyTorch, so don't mark this as a duplicate of <a href="https://stackoverflow.com/questions/63603692/pytorch-how-to-do-inference-in-batches-inference-in-parallel">this question</a>.)</p>
  <p>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:</p>
  <pre><code>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()
  </code></pre>
  <p>This crashes with the error message:</p>
  <pre><code>AttributeError: Can't pickle local object 'Agent.learn.&lt;locals&gt;.spawned'
  </code></pre>
  <p>So how do I make this work?</p>
  <p>(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.)</p>