Unable to convert the final mul, add, topk nodes to HLS

[arishsatheesan]

Hi,

I am trying to implement a custom model using FINN.
This is the model:

class QuantWeightActLeNet(Module):
def init(self):
super(QuantWeightActLeNet, self).init()
self.quant_inp = qnn.QuantIdentity(bit_width=4)
self.conv1 = qnn.QuantConv2d(3, 6, 5, bias=True, weight_bit_width=weight_bit_width)
self.relu1 = qnn.QuantReLU(bit_width=act_bit_width)
self.conv2 = qnn.QuantConv2d(6, 16, 5, bias=True, weight_bit_width=weight_bit_width)
self.relu2 = qnn.QuantReLU(bit_width=act_bit_width)
self.fc1 = qnn.QuantLinear(1655, 120, bias=True, weight_bit_width=weight_bit_width)
self.relu3 = qnn.QuantReLU(bit_width=act_bit_width)
self.fc2 = qnn.QuantLinear(120, 84, bias=True, weight_bit_width=weight_bit_width)
self.relu4 = qnn.QuantReLU(bit_width=act_bit_width)
self.fc3 = qnn.QuantLinear(84, 10, bias=True)

def forward(self, x):
    out = self.quant_inp(x)
    out = self.relu1(self.conv1(out))
    out = F.max_pool2d(out, 2)
    out = self.relu2(self.conv2(out))
    out = F.max_pool2d(out, 2)
    out = out.reshape(out.shape[0], -1)
    out = self.relu3(self.fc1(out))
    out = self.relu4(self.fc2(out))
    out = self.fc3(out)
    return out

model = QuantWeightActLeNet()

While performing streamlining, the final nodes are not getting converted to HLS. I am new to FINN and I tried correcting it as per some of the previous posts, but no luck.

The streamlining code I used is:

model = ModelWrapper("LeNet_WeightAct_tidy_prepost.onnx")
model = model.transform(MoveScalarLinearPastInvariants())
model = model.transform(Streamline())
model = model.transform(LowerConvsToMatMul())
model = model.transform(MakeMaxPoolNHWC())
model = model.transform(absorb.AbsorbTransposeIntoMultiThreshold())
model = model.transform(Streamline())
model = model.transform(MakeMaxPoolNHWC())
model = model.transform(absorb.AbsorbTransposeIntoMultiThreshold())
model = model.transform(Streamline())

model = model.transform(absorb.AbsorbScalarMulAddIntoTopK())
model = model.transform(RoundAndClipThresholds())

model = model.transform(InferDataLayouts())
model = model.transform(RemoveUnusedTensors())
model.save("LeNet_WeightAct_tidy_streamlined.onnx")

[iksnagreb]

The issue is likely the final Add, preceding the TopK, this is probably the bias of the last linear layer, which seems not to be followed by a quantizer. Normally, streamlining would absorb such a bias into the next MultiThreshold, but without an activation quantizer, there is no MultiThreshold. The AbsorbScalarMulAddIntoTopK transformation works, as the name suggests, only for scalar Adds and Muls and thus does not apply here, as the bias vector is not a scalar. Without streamlining-away the Add, the Mul cannot be absorbed into the TopK as well, leaving floating-point operations in the graph (everything following that Mul node), which cannot be implemented in HLS (or RTL) for now.

You could try two options now: Add an activation quantizer following you last linear layer, e.g., a QuantIdentity or a QuantReLU - it probably depends on your use-case or loss function what make sense - or disable the bias of this last layer, then there should not be any non-scalar Add preventing the Mul to be streamlined.

[arishsatheesan]

Thanks! Disabling the bias of the last layer worked.