Skip to content

Conversation

@tqchen
Copy link
Member

@tqchen tqchen commented Jun 13, 2019

Previously relay's python pass decorator(transform.module_pass, transform.function_pass) decorates functions.

However, one typical use-case of the customized pass is parameterized pipelines.
The example below is one made-up use-case when a user wants to customize the behavior of the pipeline by setting a customization flag(enable_fold).

# object style parameterized pass.
@relay.transform.module_pass
class CustomPipeline:
    def __init__(self, enable_fold):
        self.enable_fold = enable_fold
        self.cse = relay.transform.EliminateCommonSubexpr()
        self.const_fold = relay.transform.FoldConstant()

    def transform_module(self, mod, ctx):
          mod = self.cse(mod, ctx)
          if self.enable_fold:
              mod = self.const_fold(mod, ctx)
          return mod

This PR enables this kind of decoration. That decorate CustomPipeline into a new class that is a subclass of ModulePass, and enables its usage as a ModulePass.

Note that, however, we can certainly reach a similar solution using functional style programming along with the old API(see the code below).

# funcitonal style parameterized pass.
def CustomizePipeline(enable_fold):
      cse = relay.transform.EliminateCommonSubexpr()
      const_fold = relay.transform.FoldConstant()
      def transform_module(mod, ctx)
            mod = cse(mod, ctx)
            if enable_fold:
                mod = const_fold(mod, ctx)
            return mod
     return relay.transform.module_pass(transform_module)

So it was really a debate about which style that a user might prefer(FP vs objects) and what will happen as the pipeline gets more complicated. By talking to folks, the feedback that I get so far is:

  • If the parameterization behavior is simple, functional style seems to be more succinct
  • If there are a lot of local states, member functions, as the pipeline becomes more complicated, class style might be preferred.
  • One notable fact: most deep learning frameworks' compositional API use the class style, possibly due to the need to override multiple functions.

After considering these facts, I think it might be helpful to enable the object style of writing customized passes. So that certain customized pipelines like Quantization might take benefit from this improvement.

cc @jroesch @ZihengJiang @zhiics @eqy @MarisaKirisame @yzhliu @icemelon9

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant