算法面试准备 - 手撕系列第一期 - Softmax

算法面试准备 - 手撕系列第一期 - Softmax

Softmax原理图

Softmax原理图

Softmax实现代码 - 复杂版和简单版本(推荐简单版本)

• 方法一：循环计算

import torchdef softmax(X): # X为Tensor向量，大小为(batch_size,len)# 方法一，复杂版本for i in range(X.size()[0]):# 取出某行的Tensor# 为了防止数据的上下，先把数据减去最大值X[i]-=max(X[i].clone())X[i]=torch.exp(X[i])X[i]/=X[i].sum()return Xif __name__=='__main__':X=[[i for i in range(4)] ,[j for j in range(1,5) ],[5,4,3,2]]X=torch.tensor(X,dtype=torch.float)print(softmax(X))

• 方法二：广播利用torch函数

import torchdef softmax(X): # X为Tensor向量，大小为(batch_size,len)# 方法二：简单版本# 增加一步，防止数据上下溢出# (batch_size, 1)X_max,X_index=torch.max(X,dim=1,keepdim=True) # 让其保持二维X -= X_max# 取expX_exp = torch.exp(X)# 求和从1维求和得到的是(batch_size,1)X_sum=X_exp.sum(dim=1,keepdim=True)return X_exp/X_sumif __name__=='__main__':X=[[i for i in range(4)] ,[j for j in range(1,5) ],[5,4,3,2]]X=torch.tensor(X,dtype=torch.float)print(softmax(X))

参考

小红书资料：哲学专业转算法岗｜手撕softmax - 小红书