torch.jit support

[mdraw]

All network layers in this repo can be traced with torch.jit.trace(), but their control flow won't work correctly in traced modules (e.g. training/eval mode is not respected and the check in

dropblock/dropblock/scheduler.py

Line 16 in 16a518a

if self.i < len(self.drop_values):

will be ignored if it's true during the tracing process), unlike PyTorch's native Dropout implementations. I've tried to fix this by rewriting the control flow explicitly in TorchScript, but I didn't get it to work yet.

Since tracing this code does not necessarily emit warnings (see example below), I think this incompatibility should be documented here to make sure no one mistakenly trains traced networks. OTOH networks that are traced in eval mode after a complete training should work as intended, as long as they are not put in training mode again.

Example code that produces a wrong TracedModule without any warning (using PyTorch 1.0.0):

import torch
from dropblock import DropBlock2D, LinearScheduler

drop_block = LinearScheduler(
    DropBlock2D(block_size=3, drop_prob=0.),
    start_value=0.,
    stop_value=0.25,
    nr_steps=5
)

x = torch.randn(1, 1, 8, 8)
traced = torch.jit.trace(drop_block, x)

[miguelvr]

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.

[miguelvr]

@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.

[miguelvr]

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.

[mdraw]

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.

[miguelvr]

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

[miguelvr]

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