Skip to content
Advertisement

How to get the last index of model’s prediction?

I am new to PyTorch. I have a variable pred which has a list of a tensor.

JavaScript

So I wanted to access the last element which is the class. I did that by first converting the list into a tensor.

JavaScript

Now, how do I access the last element or is there any better/efficient way of doing this?

EDIT: For further reference here’s the code that performs classification task.

JavaScript

During prediction pred consists of x1, y1, x2, y2, conf, and class.

e.g. pred = [tensor([[176.64380, 193.86154, 273.84702, 306.30405, 0.83492, 2.00000]])]

If there are no predictions made by the model then pred is simply empty.

e.g. pred = [tensor([], size=(0, 6))]

Presently my program stops prediction if it receives an empty tensor and throws an error:

JavaScript

EDIT1: It seems to work when I check length of the pred but I get this error when there are two or more rows in tensor.

JavaScript

How do I make my program sort of just ignore if there are no predictions made at a certain frame and continue onto the next frame?

Advertisement

Answer

You can select the last element with the index notation on the 3rd axis, then broadcasting to a 1D tensor:

JavaScript

However, I would rather use torch.cat on pred, this avoids creating a new axis:

JavaScript

Edit – you may check if the tensor is empty beforehand with:

JavaScript
User contributions licensed under: CC BY-SA
3 People found this is helpful
Advertisement