Issue Converting Network to Hardware Layers with FINN

[PedroAntunes178]

Hi everyone,

I'm encountering an issue (cycle-free graph violated: partition depends on itself) while converting the following network to hardware layers using the FINN dataflow compiler:

import torch
import torch.nn as nn
import torch.nn.functional as F

import brevitas.nn as qnn

class CNetPlusScalar(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant_inp_img = qnn.QuantIdentity(return_quant_tensor=True)
        self.quant_inp_scalar = qnn.QuantIdentity(return_quant_tensor=True)
        self.conv1 = qnn.QuantConv2d(1, 6, 5, bias=True, return_quant_tensor=True) 
        self.conv2 = qnn.QuantConv2d(6, 16, 5, bias=True, return_quant_tensor=True) 
        self.conv3 = qnn.QuantConv2d(16, 32, 5, bias=True, return_quant_tensor=True) 
        self.conv4 = qnn.QuantConv2d(32, 32, 5, bias=True, return_quant_tensor=True)
        self.fc1 = qnn.QuantLinear(32 * 28 * 28, 120, bias=True, return_quant_tensor=True)
        self.fc2 = qnn.QuantLinear(120, 84, bias=True, return_quant_tensor=True)
        self.fc3 = qnn.QuantLinear(84+1, 1, bias=True, return_quant_tensor=False)
        # Separate pooling instances for each stage
        self.pool1 = nn.MaxPool2d(2, 2)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.pool3 = nn.MaxPool2d(2, 2)
        self.pool4 = nn.MaxPool2d(2, 2)
        # Instantiate separate activation layers
        self.act_conv1 = qnn.QuantReLU(return_quant_tensor=True)
        self.act_conv2 = qnn.QuantReLU(return_quant_tensor=True)
        self.act_conv3 = qnn.QuantReLU(return_quant_tensor=True)
        self.act_conv4 = qnn.QuantReLU(return_quant_tensor=True)
        self.act_fc1   = qnn.QuantReLU(return_quant_tensor=True)
        self.act_fc2   = qnn.QuantReLU(return_quant_tensor=True)

    def forward(self, image, background_scalar):
        image = self.quant_inp_img(image)
        x = self.pool1(self.act_conv1(self.conv1(image))) 
        x = self.pool2(self.act_conv2(self.conv2(x))) 
        x = self.pool3(self.act_conv3(self.conv3(x))) 
        x = self.pool4(self.act_conv4(self.conv4(x))) 
        x = torch.flatten(x, 1) 
        x = self.act_fc1(self.fc1(x))
        x = self.act_fc2(self.fc2(x))
        x = self.quant_inp_scalar(x)
        background_scalar = self.quant_inp_scalar(background_scalar)
        x = torch.cat((x, background_scalar), dim=1)
        x = self.fc3(x)
        return x

model = CNetPlusScalar()

I’ve attached part of the step_convert_to_hw.onnx graph. As you can see, the last pooling layer does not have the fpgadataflow modifier.

Does anyone know what custom steps might be required or what I might be doing wrong? Any guidance would be greatly appreciated!

Thanks in advance!