Hi. I encountered an issue when trying to convert a PyTorch FCN model into Relax IRModule following the tutorial Import from existing models: torch.nn.dropout seems to be discarded.
In Relay IR, torch.nn.Dropout is converted to relay.nn.dropout. However, in Relax, though relax.op.nn.dropout is defined, torch.nn.Dropout is not converted to it but discarded.
Actually I found a piece of code in tests/python/relax/test_frontend_from_fx.py:
# dropout
class Dropout1(Module):
def __init__(self):
super().__init__()
self.dropout = torch.nn.Dropout(0.5)
def forward(self, input):
return self.dropout(input)
class Dropout2(Module):
def forward(self, input):
return torch.dropout(input, 0.5, train=True)
@tvm.script.ir_module
class expected_dropout:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32")
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1
R.output(gv)
return gv
verify_model(Dropout1(), input_info, {}, expected_dropout)
verify_model(Dropout2(), input_info, {}, expected_dropout)
It seems that the specific behaviour is known to TVM group. I’m wondering two questions: Was this change made intentionally, meaning that TVM now only focuses on model inference? Is there any way to reserve these training-related operators or model structures?