Hi @masahi, thanks for your answer. I’m doing QAT with Brevitas library.
Thank you also for your answer @cgerum. I’m already exporting to TorchScript first. Actually, the tutorial of TVMCon2021 does not go far enough to reproduce the issue, because the issue appears after the Brevitas export:
relay.frontend.from_pytorch(traced_model, ["input", img_size])
The code on the readme of the Brevitas GitHub page (https://github.com/Xilinx/brevitas) also raises the same error. I already mentioned it 2 times on the Brevitas Gitter page but I have no answer.
Here is an easy code to reproduce the issue:
import torch
from tvm import relay
import brevitas.nn as qnn
from brevitas.quant.scaled_int import Int8ActPerTensorFloat, Uint8ActPerTensorFloat, Int8WeightPerTensorFloat, Int8Bias
from brevitas.export import export_pytorch_quant
class BasicQNN(torch.nn.Module):
def __init__(self):
super(BasicQNN, self).__init__()
self.quant_img = qnn.QuantIdentity(act_quant=Uint8ActPerTensorFloat, return_quant_tensor=True)
self.conv1 = torch.nn.Sequential(
qnn.QuantConv2d(in_channels=1, out_channels=16, kernel_size=3, stride=2, padding=1, bias=False,
weight_quant=Int8WeightPerTensorFloat, output_quant=Int8ActPerTensorFloat,
return_quant_tensor=True),
torch.nn.ReLU()
)
self.conv2 = torch.nn.Sequential(
qnn.QuantConv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1, bias=False,
weight_quant=Int8WeightPerTensorFloat, output_quant=Int8ActPerTensorFloat,
return_quant_tensor=True),
torch.nn.ReLU()
)
self.fc = torch.nn.Sequential(
qnn.QuantLinear(in_features=32*7*7, out_features=10, bias=False,
weight_quant=Int8WeightPerTensorFloat, output_quant=Int8ActPerTensorFloat,
return_quant_tensor=False)
)
def forward(self, x):
x = self.quant_img(x)
x = self.conv1(x)
x = self.conv2(x)
x = x.reshape(x.shape[0], -1)
x = self.fc(x)
return x
def main():
# Create pytorch model
model = BasicQNN().cpu().eval()
img_size = (1, 1, 28, 28)
# Export to torchscript
traced_model = export_pytorch_quant(model, input_shape=img_size)
# Export to TVM
mod, params = relay.frontend.from_pytorch(traced_model, ["input", img_size])
if __name__ == "__main__":
main()