在 PyTorch 中

可以使用 F.softmax 函數來計算 softmax

用於將一個具有多個元素的向量轉換為一個概率分佈

其中每個元素都介於 0 和 1 之間，並且總和為 1

以下是範例：

 

import torch
import torch.nn.functional as F

# 假設你有一個具有多個元素的一維向量
x = torch.tensor([1.0, 2.0, 3.0])

# 對向量 x 執行 softmax 轉換
softmax_output = F.softmax(x, dim=0) # 這邊是1維才等於0喔

print(softmax_output) #輸出 tensor([0.0900, 0.2447, 0.6652])

 

dim=0 代表 softmax 將在第一個維度進行計算

輸出總和等於 1的結果

所以如果輸入是2維

那這邊 dim 就要等於 1

給大家參考囉