Blur image directly after inference on GPU

Is there any way to run blurring on a torch tensor on a GPU? I am running a segmentation model and would like to blur the segmented areas of the original image directly on the gpu afterwards (not put it onto the CPU first to blur).

Does torchvision.transforms.GaussianBlur do what you need?

import torch, torchvision
device = torch.device("cuda")
tensor = torch.rand(3, 224, 224).to(device)
blurred_tensor = torchvision.transforms.GaussianBlur(3)(tensor)


Cool thanks you! Is the blurring operation also done on the gpu?

I believe so. Perhaps @ptrblck can you confirm?

1 Like

Yes, you are right and internally this method will be used, which will apply a conv2d on the input tensor using the specified img.device arguments.
If a PIL.Image is passed it’ll be transformed to a tensor before the blur transformation is applied and the result will then be transformed back to a PIL.Image as seen here (so the CPU will be used in this case).

1 Like

Cool thank you very much!