深度学习一 —— 手撕softmax

s o f t m a x ( x i ) = e x i ∑ j ( e x j ) softmax(x_i) = frac{e^{x_i}}{sum_j(e^{x_j})} softmax(xi​)=∑j​(exj​)exi​​

import numpy as np
 
def softmax(x, axis = 1):
    assert (len(x.shape) > 1, "dimension must be larger than 1")
    x -= np.max(x, axis=axis, keepdims=True)
    x = np.exp(x) / np.sum(np.exp(x), axis=axis, keepdims=True)
    return x

注意：

为了稳定的计算softmax概率，防止 e x e^x ex过大，出现 n a n nan nan的情况，会选择减去其最大值。