[jit] Support for modules that have hooks when compiling

[rami3e]

🚀 Feature

Currently forward or backward hooks (i.e concatenation of features in UNET) are not supported in jit tracing.

Run into this error when attempting: ValueError: Modules that have hooks assigned can't be compiled

source: torch.jit.TracedModule(ScriptModule)

[suo]

Hi, could you explain a little bit more about your use case? For example, would it be possible as a workaround to inline the hook into the forward() of your model?

[rami3e]

Sure. It is similar to U-net style skip connections but I also pass these features through another conv layer. It is actually in the forward() currently. ie.

def forward(self, input, enc_feats):  
     # ....
     y2 = torch.cat((y0, (nn.Conv2d...)))

[rami3e]

has this been addressed elsewhere ?

[eugeneware]

@suo It's for getting U-nets to work where hooks are used to store activations from the forward pass. See code here for the model that many people in the fast.ai course are trying to compile.

It turns out that if you just comment out that error the compilation works.

So it would be great if there was some option that we could pass to bypass that exception, and perhaps just get a warning.

You can follow some of the discussion in this forum thread here for some more context.

Thanks for your great work!

[eugeneware]

FYI this issue is related to this issue also.

[suo]

@zdevito do you remember why we explicitly disallow hooks?

[zdevito]

I can't remember. We should revisit it and see if wean just enable it.

[pshashk]

nn.utils.spectral_norm is also implemented via hook.

torch.jit.trace(nn.utils.spectral_norm(nn.Linear(8,8)), torch.rand(4,8))

ValueError: Modules that have hooks assigned can't be compiled

[fuzic]

weight norm also. Are models using these just unexportable?

https://pytorch.org/docs/stable/_modules/torch/nn/utils/weight_norm.html

[eellison]

@zdevito @suo any thoughts ?

[eugeneware]

Thank you @eellison ! 🎉 👏