Are MONAI transforms differentiable?

[joaolcguerreiro]

Hi! I would like to know if MONAI's transforms are differentiable. For example if I have a network that generates images. I need to transform the network output before passing it to the criteria and optimise the network. If I use for example CenterSpatialCrop, will it break the computational graph in pytorch? Or will the gradients flow back to the network?

Thank you :)

[joaolcguerreiro]

I will explain better here. Imagine I have a neural network with goal of generating realistic images. One of the loss components is a perceptual loss. In a perceptual loss I compare the features extracted from the network output with the features extracted from the ground-truths. For this feature extraction I use a pre-trained VGG19 network (its layers are freeze). Before I compute the features I need to transform my generates images and ground-truth to a VGG input compatible shape... therefore I need to transform them beforehand using CenterSpatialCrop. If I then pass the transformed images to the criteria (used to compute the perceptual loss) and then call loss.backwards() will the computational graph work well and the gradients from the generated images will flow back to the network? Or since I transformed the generated images the computational graph was broken?

output_img = network(input_img)

transforms = monai.transforms.Compose([NormalizeIntensity(mean=mean, std=std), transforms.Resize(spatial_size=(256, 256, 256), mode="trilinear"), transforms.CenterSpatialCrop(roi_size=(224, 224, 224))])

transformed_output_img = transforms(output_img)
transformed_ground_truth_img = transforms(ground_truth_img)

loss = criterion(transformed_output_img, transformed_ground_truth)

# Backward pass
loss.backward()
network_optimizer.step()

Essentially, will the gradients still flow back to the network G if I pass to the criteria a transformed network output? Like if I use a MONAI transform before passing it to the criteria will the gradients flow back to the network?

[joaolcguerreiro]

Okay, I have transversed the entire computational graph and I can see nodes from the network. This means the computational graph was not broken after the transforms, i.e., MONAI's transforms seem to be differentiable.

In case someone need I can provide the function I used to transverse the graph :)

[wyli]

Thanks for the requests and comments. The current status of the transforms is that we try to keep the gradients as much as possible when using the pytorch APIs. Certain transforms and their options may trigger numpy/cupy backends, and those gradient wont be preserved.

[joaolcguerreiro]

I checked Resize and CenterSpatialCrop and they seem to be differentiable. Do you know which transforms are not?