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

Question

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 1

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

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

说明：在pytorch中，可以直接使用prob>0.5获取torch.bool类型张量。然后你可以通过.float().

Answer 2

为什么不考虑使用无循环解决方案？也许像下面这样就足够了：

In [34]: output = torch.tensor([0.4481, 0.4014, 0.5820, 0.2877]) 

# subtract off the threshold value (0.5), create a boolean mask, 
# and then cast the resultant tensor to an `int` type
In [35]: result = torch.as_tensor((output - 0.5) > 0, dtype=torch.int32) 

In [36]: result        
Out[36]: tensor([0, 0, 1, 0], dtype=torch.int32)