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.

print(pred)
output: [tensor([[176.64380, 193.86154, 273.84702, 306.30405, 0.83492, 2.00000]])]

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

x = torch.stack(pred)
output: tensor([[[176.64380, 193.86154, 273.84702, 306.30405, 0.83492, 2.00000]]])

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.

def classify_face(image):
    device = torch.device("cpu")
    img = process_image(image)
    print('Image processed')
    # img = image.unsqueeze_(0)
    # img = image.float()
    
    pred = model(img)[0]
    
    # Apply NMS
    pred = non_max_suppression(pred, 0.4, 0.5, classes = [0, 1, 2], agnostic = None )
    if classify:
        pred = apply_classifier(pred, modelc, img, im0s)
    #print(pred)
    
    model.eval()
    model.cpu()
    
    print(pred)
        
    # output = non_max_suppression(output, 0.4, 0.5, classes = class_names, agnostic = False)
        
    #_, predicted = torch.max(output[0], 1)
    #print(predicted.data[0], "predicted")

    classification = torch.cat(pred)[:, -1]
    index = int(classification)
    print(names[index])
    return names[index]

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:

Traceback (most recent call last):
  File "WEBCAM_DETECT.py", line 168, in <module>
    label = classify_face(frame)
  File "WEBCAM_DETECT.py", line 150, in classify_face
    index = int(classification)
ValueError: only one element tensors can be converted to Python scalars

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.

[tensor([[212.38568, 117.47020, 339.35773, 266.00513,   0.74144,   2.00000],
        [214.60651, 118.50694, 339.90192, 265.91696,   0.44277,   0.00000]])]
#################
#################
Traceback (most recent call last):
  File "WEBCAM_DETECT.py", line 172, in <module>
    label = classify_face(frame)
  File "WEBCAM_DETECT.py", line 154, in classify_face
    index = int(classification)
ValueError: only one element tensors can be converted to Python scalars

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?

Answer

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

x[:, :, -1].view(-1)

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

torch.cat(pred)[:, -1]

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

if len(pred) == 0:
    return None

Advertisement

Answer