torch.jit support about dropblock HOT 6 CLOSED

I haven't messed with PyTorch 1.0 tracing features yet. However, this is something that interests me and that I would like very much to support.

Both modules will have indeed to be converted to TorchScript due to control flow.

I'll have a look today to see if I can get any progress.

from dropblock.

@mdraw I'm surprised you had no errors tracing the DropBlock2D module. I'm converting it to ScriptModule and it requires a bunch of subtle changes in the code.

from dropblock.

I think it might not be possible to trace the LinearScheduler because it changes the value of the drop_prob of a module, which I think can not be done with TorchScript.

I have to investigate this, however.

from dropblock.

I'm surprised you had no errors tracing the DropBlock2D module.

In my (limited) experience, implementing things as ScriptModules is currently pretty hard because of the limited language and library features that the compiler understands, but tracing has somehow worked even for complicated models - although you have to watch out for code that the tracer doesn't recognize.

I think it might not be possible to trace the LinearScheduler because it changes the value of the drop_prob of a module, which I think can not be done with TorchScript

Yes, I also don't know how this could be done. __constants__ work fine for constants, but I don't see any way to define module-level variables that you can change dynamically. I just checked: torch.nn.Dropout2d doesn't support changing p either if it's in ScriptModule form.

Besides this issue, your implementation in #27 looks good to me. DropBlock2D and DropBlock3D work as intended now in my local tests.

from dropblock.

Yes, the DropBlock modules are now working fine. I didn't merge yet due to the LinearScheduler. I will do some more research but I really think it is not possible at the moment

from dropblock.

closing this as DropBlock does not benefit from JIT scripting, as it can be traced for inference time without problems.

from dropblock.