MONAI Transforms inherit torch.nn.Module

[bhashemian]

Is your feature request related to a problem? Please describe.
MONAI array transforms are all derived from Transform base class by implementing __call__ method, and similar to torchvision transform can be chained using Compose. However, in torchvision transforms without losing any backward compatibility are inherited fromnn.Module (since v0.8.0). This has some obvious benefits:

1. They can be used torchscripted with torch.jit.script
2. They can be chained using nn.Sequential. This will allow the transforms be chained with the model as well and let them run on the same device that model is being run.

Describe the solution you'd like

1. Derive all transform from torch.nn.Module instead of Transform
2. Implement forward for all transforms instead of __call__
3. Make them work with torch.Tensors (Default to using Pytorch APIs for monai.transforms (ETA TBD) #2231)

[wyli]

torchvision only supports torchscript for some of the transforms, and the current version of torchscript has its limitations, so I think "deriving all transforms" from torch.nn.Module needs more motivations at the moment. We can add new extension class of nn.Modules to monai.networks, and the user can optionally use them in the preprocessing directly.

[vfdev-5]

My 2 cents for this discussion. One of the goals for torchvision.transforms is to support torch.Tensor as input in addition to PIL image, such that transforms could be executed on other accelerators, scripted (and thus be a part of the models forward)

Torchvision's transformations not derived from torch.nn.Module are mostly PIL/numpy related or kept for BC as users can do a lot of unexpected things with classes like Compose.

[rijobro]

Interesting. Would be fine if all transforms natively support torch.Tensor as input. However, we'll have a few transforms that will need to convert to numpy at the start of the __call__ method, do we think this would cause any problems if these transforms also inherited from torch.nn.Module?

[vfdev-5]

I think torch tensor support for transforms can be dissociated from a need to inherit from torch.nn.Module.

Technically, there is no major blockers to make a transform class to inherit from torch.nn.Module. Then the question is what we would like to get with that.

Main blocker for torchvision to rewrite Compose inherited from torch.nn.Module , as far as I can say, is/was support of any callable or lambda transformations which could not be torchscripted or simply inserted into nn.Sequential.

However, we'll have a few transforms that will need to convert to numpy at the start of the call method, do we think this would cause any problems if these transforms also inherited from torch.nn.Module?

I think torchscript wont work with numpy datatype for these transformations even if inherited from torch.nn.Module.