Search code examples
pythonpytorchtorchvision

PyTorch [1 if x > 0.5 else 0 for x in outputs ] with tensors

I have a list outputs from a sigmoid function as a tensor in PyTorch

E.g

output (type) = torch.Size([4]) tensor([0.4481, 0.4014, 0.5820, 0.2877], device='cuda:0',

As I'm doing binary classification I want to turn all values bellow 0.5 to 0 and above 0.5 to 1.

Traditionally with a NumPy array you can use list iterators:

output_prediction = [1 if x > 0.5 else 0 for x in outputs ]

This would work, however I have to later convert output_prediction back to a tensor to use

torch.sum(ouput_prediction == labels.data)

Where labels.data is a binary tensor of labels.

Is there a way to use list iterators with tensors?

Solution 
  • prob = torch.tensor([0.3,0.4,0.6,0.7])

out = (prob>0.5).float()
# tensor([0.,0.,1.,1.])

    Explanation: In pytorch, you can directly use prob>0.5 to get a torch.bool type tensor. Then you can convert to float type via .float().