[feature request] Derive all transforms classes from nn.Module?

[vadimkantorov]

🚀 The feature

Currently it's already the case for most transform classes, except T.Compose / T.RandomChoice (these could probably be derived from nn.Sequential / nn.ModuleList). This also would not break any compat, and improve consistency and composability (sometimes only nn.Modules are accepted as parts of other existing models)

Motivation, pitch

N/A

Alternatives

No response

Additional context

No response

cc @vfdev-5 @datumbox @bjuncek

[datumbox]

@vadimkantorov You are right to say that. There is a handful of them that don't derive from it and the reason is mainly that they might not be JIT-scriptable and that might break BC. I think there is a lot of merit to make all classes inherit from nn.Module. The new Transforms API is the place to address it (see here).

[vadimkantorov]

T.Compose probably is scriptable, the rest probably too if we migrate from random. to using torch random functions? But yeah, that'd be a switch to a different random algo

[vfdev-5]

@vadimkantorov have you seen discussion from this PR (#2645), why classes like T.Compose were not derived from nn.Module etc ? TLDR, keep BC with previous versions and support lambda functions.

[vadimkantorov]

Thanks for the reference! I tried to understand why deriving from nn.Sequential would existing some code, but I couldn't since the discussion is quite dispersed there.

What are concrete examples you had in mind in that PR? Is it because you worry of the code that's directly modifying internal transforms field? I didn't see a lot of code like this in the wild. Some solutions I can see:

• deprecate these direct modifications (and provide some explicit replacement methods that will auto-wrap)
• make ModuleList's auto-wrap functions into Modules or lift the constraint of all elements of _modules field to be strictly modules and allow functions / lambdas as well (related issues: Can't reassign a child module instance to a lambda pytorch#51896 [discussion] [python] Generic way for generating module-wrappers for stateless parametrized functions pytorch#51782)
• introduce a new TupleSequential ([feature request] TupleSequential module or an option for Sequential for unpacking tuples (i.e. support varargs) pytorch#55457), deprecate use of Compose and at some point in the future alias Compose as TupleSequential and deprecate Compose for good
• support [feature request] TupleSequential module or an option for Sequential for unpacking tuples (i.e. support varargs) pytorch#55457

[vfdev-5]

Thanks for the reference! I tried to understand why deriving from nn.Sequential would existing some code, but I couldn't since the discussion is quite dispersed there.

First point is that we can not do basic thing like: T.Compose([T.Resize(200), lambda x: x * 0.5, my_transform_plain_function]) as lambda functions are not nn.Module and then T.Compose being inherited from nn.Module or nn.Sequential wont accept that.
To solve this we have to wrap all callables non-nn-modules with something from torchvision.

Next, if we wrap these custom functions by something then user can lose the access to its possible attributes etc. Imagine,
t = T.Compose([T.Resize(200), MyCustomTransform(*args, **kwargs)]) and t.transforms[1].call_custom_method(...). This wont be possible if we do not do some __getattr__ magic...

[vadimkantorov]

Thanks for explaining! Now I see!

My proposal would be:

• support [discussion] [python] Generic way for generating module-wrappers for stateless parametrized functions pytorch#51782 for core
• support [feature request] TupleSequential module or an option for Sequential for unpacking tuples (i.e. support varargs) pytorch#55457 or introduce a separate TupleSequential (hopefully in core :))
• update examples / docs for old-style examples to use TupleSequential instead of Compose
• Introduce module-derived T.Lambda (maybe in core [discussion] [python] Generic way for generating module-wrappers for stateless parametrized functions pytorch#51782)
• deprecate Compose / RandomChoice used with pure callables - first warn and ask to use Module-derived T.Lambda, then start failing
• migrate to use of Module-derived TupleSequential / Sequential / RandomChoice

[vadimkantorov]

Maybe RandomChoice is more safe for conversion to nn.ModuleList, as it's not as old class

[vfdev-5]

Proto transforms now all derive from nn.Module

[vadimkantorov]

Is there now a nn.Sequential/nn.ModuleList-derived TupleSequential? Or what is the format of instances in proto? Are they now all dicts? Is there now also a nn.Module-derived T.Lambda wrapper?