PyTorch [1, если x> 0,5, иначе 0 для x в выходных данных] с тензорами

Question

У меня есть список выводов из сигмовидной функции в качестве тензора в PyTorch

Например:

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

Поскольку я делаю двоичную классификацию, я хочу повернуть все значения ниже 0,5 до 0и выше от 0,5 до 1.

Традиционно с массивом NumPy вы можете использовать итераторы списка:

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

Это будет работать, однако мне придется позже преобразовать output_prediction обратно в тензор, чтобы использовать

torch.sum(ouput_prediction == labels.data)

Где label.data - двоичный тензор меток.

Есть ли способ использовать итераторы списка с тензорами?

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

Почему бы не подумать об использовании решения loop less?Может быть, что-то вроде ниже будет достаточно:

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)