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?
Answer
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().