TOpPattern has not been registered for nn.dropout

Hi there, I am working on relay gradients operation and trying to feed the bacward graph into autoscheduler to search. However, I meet errors TOpPattern has not been registered for nn.dropout when the DAG contains backward operations. Any advise about how to fix the issue?

model = nn.Sequential(
    nn.Conv2d(3, 3, kernel_size=3, padding=1),
    nn.BatchNorm2d(3),
    nn.Dropout()
)

# We grab the TorchScripted model via tracing
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()
# scripted_model
input_name = "input0"
shape_list = [(input_name, input_data.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)

# dummy gradients as an placeholder
@register_gradient("nn.dropout")
def dropout_grad(orig, grad):
    return [zeros_like(_) for _ in orig.args]

mod = relay.transform.InferType()(mod)
bwd_mod = relay.transform.gradient(mod['main'], mode="first_order")

tgt_platform = "llvm"
target = tvm.target.Target(tgt_platform)
tasks, task_weights = auto_scheduler.extract_tasks(bwd_mod, None, target)
"""
Get errors with GraphExecutorCodegen for task extraction. Fallback to VMCompiler.
TVMError: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (idx < data_.size() && data_[idx].second != 0) is false: Attribute TOpPattern has not been registered for nn.batch_norm

"""

IIRC Dropout isnā€™t really supported, when a model is run it is removed here https://github.com/apache/tvm/blob/main/src/relay/transforms/simplify_inference.cc#L201.

You can add support for dropout, there is a random number generator somewhere in tvm.

1 Like

Hi Andrew,

Thanks for sharing the pointer. But I donā€™t think I called GraphSimplifier in my code and Dropout should be kept as it is as shown in the IR. Does the relay.transform.gradient automatically call simplifier?

Ah yes, I donā€™t believe transform.gradient will call the simplify inference I pointed to, but my point is that dropout has no implementation and if you have a model with dropout you will get an error like the one you see.

Therefore the way to run models with dropout is by removing it with the pass I linked (which is what is supposed to done in inference time anyway).

For your use case however I believe you need to define the implementation first and then the gradient.

1 Like

@altanh may have something to say about dropout.

Thanks Andrew, so your point is that the Dropout actually has no implementation (either fwd and bwd) and in current execution this is automatically removed by certain transformations?

If that is the case, can you advise how can I add an implementation to the relay? I would like to make PR to support this.

Yep, thatā€™s right.

I think this Adding an Operator to Relay ā€” tvm 0.8.dev0 documentation will be helpful, specifically step 5.

Now you donā€™t have a compute implementation so you need to create this. Unfortunately I donā€™t know a good tutorial to do this and the best way might be to just look at the code like of the op in the tutorial above ^ :confused: .

https://github.com/apache/tvm/blob/main/python/tvm/topi/random/kernel.py ā† has a source of randomness for dropout you can use for this.

Finally, I think @altanh might have done this before. He might be more active on the discord than here.

1 Like

Hi Andrew,

One more question, is similar behavior also made for the bn layer? Says it was fused to conv layer and does not have an actual implementation.

Yes that is correct. Though I believe someone was planning to work on this one in the next week.

1 Like