Unable to run model only in DLA (without GPU)

Description

We are experimenting with DLA in our Jetson Xavier Embedded Device.

We first ran YOLOV5 model in DLA using following steps.

  • We trained a YOLOV5 model in Host Machine
  • We then converted the model to ONNX
  • We transferred this ONNX file to Jetson Xavier machine and used
trtexec --onnx=model_gn.onnx --shapes=input:1x3x640x640 --saveEngine=model_gn.engine --exportProfile=model_gn.json --int8 --useDLACore=0 --allowGPUFallback --useSpinWait --separateProfileRun > model_gn.log

This perfectly worked as expected.

Now we wanted to do one more experiment, where we wanted to run model only in DLA (basically we wanted to avoid using GPU). However, we knew that, DLA supports only few types of layers. So, we build a very small network, which is below,

class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size) 
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)  
    
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

We followed the same steps described above for YOLOV5 but without --allowGPUFallback. We are getting error saying fc1.weight not supported.

Can you please tell, where I am going wrong?

Hi @aravind.chakravarti ,
Can you please help us with the model to repro the issue.

Thanks

Hi @AakankshaS

Here is the code (build on Colab)

It has up to ONNX export.

Once, I export the ONNX from Colab, I convert the ONNX to DLA using below code in AGX Jetson Xavier

trtexec --onnx=MNIST_Classifier.onnx --shapes=input:1x784 --saveEngine=MNIST_Classifier.engine --exportProfile=MNIST_Classifier.json --int8 --useDLACore=0 --useSpinWait --separateProfileRun > model_gn.log

The images, I used are here:
https://drive.google.com/drive/folders/1maNOy9s6Kfyt2xydfV7zRfnPxn2EJDgN?usp=sharing

Regards,
Aravind

Hi @AakankshaS,

Do you need any other information?

Thanks,
Aravind

Hi @AakankshaS

Basically I want to run a small network (even a 3 layer network is fine) in Jetson DLA. Do you have any suggestions?

Regards,
Aravind

Hi @AakankshaS

Any update? Sorry for troubling you. I have tried another network which is much simpler than any thing before.

Even this code is not exporting to engine file without --allowGPUFallback