I am currently implementing this paper (https://arxiv.org/abs/1710.02410 1) in PyTorch and I have ran into an issue training the model.
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.
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.
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:
	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)

Torch.multiprocessing didn’t work out. Currently I am trying to use pathos for multiprocessing. I am getting an error however:
TypeError: can't pickle SwigPyObject objects

So I am pretty stuck lol. Is there a elegant way of doing this that I am missing? any ideas?
Thanks!!