Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RELAY][PASS] Enable decorating python class as Pass #3364

Merged
merged 1 commit into from
Jun 17, 2019

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