Thanks – the copy in the extract_mask function moves the current batch of images from CPU to a GPU device. Here is the pseudo code for the function:
def extract_mask(self, image_batch):
     #image_batch is the batch obtained from the dataloader.
     image_var = Variable(torch.Tensor(len(image_batch), 3, 224, 224), requires_grad=False).type(self.dtype_float).cuda(device=self.gpu_id)
     image_var.data.copy_(image_batch)
     output = self.model(image_var).cpu().numpy()
     return output

I probably should have mentioned this earlier – I am using pytorch0.3.1. Here’s a workaround which avoids the copy altogether.
def extract_mask(self, image_batch):
     #image_batch is the batch obtained from the dataloader.
     image_batch = Variable(torch.squeeze(image_batch), requires_grad=False).type(self.dtype_float).cuda(device=self.gpu_id)
     output = self.model(image_var).cpu().numpy()
     return output

This definitely speeds up the dataloader…

image.png3812×2126 283 KB