將 nn.Softmax 轉換為 torch.tensor

Question

我的神經網絡末端有一個 softmax function。 我希望將概率設為 torch.tensor。為此，我正在使用torch.tensor(nn.softmax(x))並出現錯誤RuntimeError: Could not infer dtype of Softmax 。

請問我在這里做錯了什么，或者有沒有其他方法可以做到這一點。

Answer 1

nn.Softmax是 class。 你可以像這樣使用它：

import torch

x = torch.tensor([10., 3., 8.])
softmax = torch.nn.Softmax(dim=0)

probs = softmax(x)

或者，您可以使用功能 API torch.nn.functional.softmax ：

import torch

x = torch.tensor([10., 3., 8.])

probs = torch.nn.functional.softmax(x, dim=0)

它們是等價的。 在這兩種情況下，您都可以檢查type(probs)是否為<class 'torch.Tensor'> 。