Multi-process inference

username: urukhai
cooked: <p>I’m looking for a way to do inference on multiple GPUs for an application where inference speed is critical.  I have an upstream process that delivers images to a vision model in batches of 50-100.  I have tried two ways of splitting the batches up so that each worker gets a different partition, but neither has been fully successful.  I also messed around with DistributedDataLoader and other ideas, but this is as far as I have gotten so far.  I would appreciate advice or suggestions.</p>
<p><strong>METHOD 1</strong> pseudo-code.<br>
This method works but only gives a 2X increase in speed with 4 GPUs and seems to spawn more processes than it should; there is a ‘print’ statement after the imports at the top of my script and it gets called 16 or more times.</p>
<pre><code class="lang-python">world_size = comm.get_world_size()
rank = comm.get_local_rank()
model = DistributedDataParallel(model, device_ids=[rank], broadcast_buffers=False)

for ib in image_batches:
	#Upstream process delivers a batch of 50 to 100 photos
	batch = get_image_batch()
	inference_dataset = DatasetFromList(batch)

	#Create a new sampler and dataloader for this particular batch
	inferencesampler = torch.utils.data.distributed.DistributedSampler(
		inference_dataset,
		num_replicas=4,
		rank=rank
		)

    testloader = torch.utils.data.DataLoader(
	    dataset=inference_dataset,
	    batch_size=bs,
	    shuffle=False,            
	    num_workers=4,
	    pin_memory=True,
	    collate_fn=trivial_batch_collator,
	    sampler=inferencesampler)     

    #Run the batch
    with torch.no_grad():
        for idx, inputs in enumerate(data_loader):
			outputs = model(inputs)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
</code></pre>
<p><strong>METHOD 2</strong> pseudo-code.  This method uses an IterableDataset in attempt to keep the dataloader outside the batch loop.  Pytorch prohibits you from re-setting the attribute values of a dataloader, which makes it impossible to switch out the dataset on the fly.  To get around that, I added a reset function to the IterableDataset so that I can pass it a new batch and reset its internals without having to change the dataloader.  The dataset works perfectly when tried in isolation (i.e., it splits the data correctly between 4 processes) but in the full program, it passes duplicate information to each worker.  I can’t figure out why (I also tried using the <code>worker_init_fn</code> method but it also failed).</p>
<pre><code class="lang-python">world_size = comm.get_world_size()
rank = comm.get_local_rank()
model = DistributedDataParallel(model, device_ids=[rank], broadcast_buffers=False)

#Instantiate the IterableDataset with dummy data outside the batch loop
inference_dataset = IterableBatchDataset(['a','b','c','d'])

testloader = torch.utils.data.DataLoader(
    dataset=inference_dataset,
    batch_size=bs,
    shuffle=False,            
    num_workers=4,
    worker_init_fn=worker_init_fn,
    collate_fn=trivial_batch_collator)     

for ib in image_batches:
	#Upstream process delivers a batch of 50 to 100 photos
	batch = get_image_batch()

	#Reset the internals of the IterableDataset
	inference_dataset.reset_dataset(batch)

	#Run the batch
	with torch.no_grad():
	    for idx, inputs in enumerate(data_loader):
			outputs = model(inputs)
	    if torch.cuda.is_available():
	        torch.cuda.synchronize()

class IterableBatchDataset(torch.utils.data.IterableDataset):
    """
    An IterableDataset for passing arbitrary-length data to a Dataloader. The
    iter() method ensures that data is not duplicated if there is more than one worker, by splitting
    the data into separate pieces and only passing each worker one piece to iterate.
    """
    def __init__(self, batchlist):
        super(IterableBatchDataset).__init__() #The init method of IterableDataset
        self.data = batchlist
        self.start = 0
        self.end = len(self.data)

    def reset_dataset(self, newdataset):
        self.data = newdataset
        self.start = 0
        self.end = len(newdataset)
        
    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        print("IterableDataset worker_info.id",worker_info.id)
        if worker_info is None:  # single-process data loading, return the full iterator
            indexes = range(self.start,self.end)
        else:  # in a worker process
            # split workload
            per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
            worker_id = worker_info.id
            iter_start = self.start + worker_id * per_worker
            iter_end = min(iter_start + per_worker, self.end)
            indexes = range(iter_start, iter_end)
        itemlist = [self.data[t] for t in indexes] 
        #print("worker_id:",worker_id,"items:",itemlist) #very useful for debugging
        return iter(itemlist)
    def __len__(self):
        return len(self.data)
    def __getitem__(self,idx):
        return self.data[idx]
</code></pre>