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
"""