softmax python 代码实现

import numpy as np
import torch

def softmax(logits):
    """
    Compute softmax values for each sets of scores in 1d
    """
    e_x = np.exp(logits - np.max(logits))
    return e_x / e_x.sum()

def softmax_torch(logits, dim=-1):
    shifted_logits = logits - logits.max(dim=dim, keepdim=True).values
    
    exp_logits = torch.exp(shifted_logits)
    return exp_logits / exp_logits.sum(dim=dim, keepdim=True)

if __name__ == "__main__":
    logits_1d = [2.0, 1.0, 0.1]
    result_1d = softmax(logits_1d)
    
    
    logits_2d = [[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]
    result_2d = softmax(logits_2d)
    
    
    logits_torch_1d = torch.tensor([2.0, 1.0, 0.1])
    result_torch_1d = softmax_torch(logits_torch_1d)
    
    logits_torch_2d = torch.tensor([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]])
    result_torch_2d = softmax_torch(logits_torch_2d, dim=-1)
    
    
    builtin_rsult = torch.nn.functional.softmax(logits_torch_2d, dim=-1)
    
    logits_3d = torch.randn([2, 3, 4])
    result_3d = softmax_torch(logits_3d, dim=-1)