Masking image on GPU

Is there a straightforward way given a mask and an image to blur the part of the image the mask corresponds to on torch.cuda (meaning doing everything on the GPU). The GaussianBlur is implemented, but how would I do the masking (would it be the same as I would do it in numpy or is there some optimized method for it)?


Yes. You can move your tensors into your device in your transformation before applying the GaussianBlur. Here is an example:

import torch
import torchvision
import torchvision.transforms as T
# ...
device = torch.device('cuda') 
transform = T.Compose([
    # ...
    T.Lambda(lambda x:
    T.GaussianBlur((7, 13))
    # ...
imgs = torch.stack([copy.deepcopy(transform(img)) for i in range(10)])

for i in range(imgs.shape[0]):
# >>> cuda:0
# ...
# >>> cuda:0

You might consider looking into this post for troubleshooting when having your data loader to have multiple workers.

I was talking rather about the masking, but I figured it out. Thanks anyway!

1 Like