How to deal with batches with a gated branch network (self-driving network)

JonathanHampton: <p>I am currently implementing this paper (<a href="https://arxiv.org/abs/1710.02410" rel="nofollow noopener">https://arxiv.org/abs/1710.02410</a>) in PyTorch and I have ran into an issue training the model.</p>
<p>The input data consists of an image, a measurement scalar, and a control scalar. The image gets passed to a convolutional module and the measurement scalar gets passed through a linear module. the outputs of these are concatenated. The final section of the model consists of several branches. The control scalar determines which branch the concatenated tensor is passed to.</p>
<p>I tested it with a batch size of 1 and it works. However my issue comes when I turn the batch size up. Since each sample in the batch goes to a (possibly different) branch and each sample in a batch is ran through the model in parallel, this causes an issue.</p>
<p>To avoid a bottleneck in the network during training i’ve been attempting to implement multiprocessing but it hasn’t worked out. Here is where im at:</p>
<pre><code class="lang-auto">	def gated_branch_function(self, j, control):
		control = int(control.item())
		""" Branches """
		if control == 2 or control == 0:
			output = self.follow_lane_branch(j)
		elif control == 3:
			output = self.left_branch(j)
		elif control == 4:
			output = self.right_branch(j)
		elif control == 5:
			output = self.straight_branch(j)
		else:
			output = self.general_branch(j)		
		return output


	def forward(self, input_data):
		''' Define variables '''
		input_image = input_data[0].permute(0,3,2,1).float()
		input_speed = input_data[1].unsqueeze(dim=1)
		control = input_data[2]

		''' Pass input data into network '''
		imageOutput = self.imageModule(input_image)
		speedOutput = self.measurementModule(input_speed)

		''' Joint sensory '''
		j = torch.cat([imageOutput, speedOutput], 1)
		j = self.jointSensoryModule(j)

		''' Branches '''
		pool = mp.ProcessPool(nodes=os.cpu_count())
		output = pool.map(self.gated_branch_function, j, control)
		return torch.stack(output)
</code></pre>
<p>Torch.multiprocessing didn’t work out. Currently I am trying to use pathos for multiprocessing. I am getting an error however:</p>
<pre><code class="lang-auto">TypeError: can't pickle SwigPyObject objects
</code></pre>
<p>So I am pretty stuck lol. Is there a elegant way of doing this that I am missing? any ideas?</p>
<p>Thanks!!</p>

JonathanHampton: <p>Update, I’ve got the following code working:</p>
<pre><code class="lang-auto">	def gated_branch_function(self, j_batch, control_batch):
		batch_outputs = []
		for j, control in zip(j_batch, control_batch):
			control = int(control.item())
			""" Branches """
			if control == 2 or control == 0:
				output = self.follow_lane_branch(j)
			elif control == 3:
				output = self.left_branch(j)
			elif control == 4:
				output = self.right_branch(j)
			elif control == 5:
				output = self.straight_branch(j)
			else:
				output = self.general_branch(j)
			batch_outputs.append(output)
		return torch.stack(batch_outputs)


	def forward(self, input_data):
		''' Define variables '''
		input_image = input_data[0].permute(0,3,2,1).float()
		input_speed = input_data[1].unsqueeze(dim=1)
		control = input_data[2]

		''' Pass input data into network '''
		imageOutput = self.imageModule(input_image)
		speedOutput = self.measurementModule(input_speed)

		''' Joint sensory '''
		j = torch.cat([imageOutput, speedOutput], 1)
		j = self.jointSensoryModule(j)

		''' Branches '''
		output = self.gated_branch_function(j, control)
		return output
</code></pre>
<p>However this results in a bottleneck, so training is slow. Is there anyway I could write gated_branch_function so that each sample in the batch gets sent to their respective branch in parallel?</p>

ptrblck: <p>You could try to add CUDA streams as described <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-streams">here</a>, but would have to take care of the synchronizations to avoid race conditions.</p>
<p>I would recommend to try it out with a “similar” but simple model and check, if you would get any performance gains at all.</p>

JonathanHampton: <p>I tested it out on a simple model and got a slight increase in performance. I rewrote the function in my model and its working:</p>
<pre><code class="lang-auto">	def gated_branch_function(self, j_batch, control_batch):
		batch_output = []
		for j, control in zip(j_batch, control_batch):
			control = int(control.item())
			s = torch.cuda.Stream()
			""" Branches """
			if control == 2 or control == 0:
				with torch.cuda.stream(s):
					output = self.follow_lane_branch(j)
					batch_output.append(output)
			elif control == 3:
				with torch.cuda.stream(s):
					output = self.left_branch(j)
					batch_output.append(output)
			elif control == 4:
				with torch.cuda.stream(s):
					output = self.right_branch(j)
					batch_output.append(output)
			elif control == 5:
				with torch.cuda.stream(s):
					output = self.straight_branch(j)
					batch_output.append(output)
			else:
				with torch.cuda.stream(s):
					output = self.general_branch(j)
					batch_output.append(output)
		torch.cuda.synchronize()
		return torch.stack(batch_output)
</code></pre>
<p>I’m guessing this is the best I can do. Still slow but at least a bit faster than before.</p>