在PyTorch分类模型中Softmax函数默认包含在CrossEntropyLoss函数中因此在分类模型的结构输出设计时无需进行Softmax激活故训练好的模型最后一层输出的是 logits未归一化的分数若需在推理时直接输出概率则需自行添加 nn.Softmax。此处可以由实现方法具体实现如下1、使用nn.Sequential将原模型与nn.Softmax封装在一起import torch import torch.nn as nn # 1. 加载已训练好的模型和权重 # 假设原模型类为 MyModel需先实例化空模型再加载权重 model MyModel() model.load_state_dict(torch.load(model_weights.pth)) model.eval() # 切换到评估模式 # 2. 封装带 Softmax 的新模型 # dim-1 表示对类别维度进行 softmax 计算 # 注意dim值需根据具体的模型识别场景进行设置 softmax_model nn.Sequential( model, nn.Softmax(dim-1) ) # 3. 保存新模型结构及权重 # 注意nn.Sequential 包装后state_dict 的 key 会变化如增加 0. 前缀 torch.save(softmax_model.state_dict(), model_with_softmax.pth) # --- 验证使用 --- # 加载时需确保结构与保存时一致 loaded_softmax_model nn.Sequential( MyModel(), # 注意dim值需根据具体的模型识别场景进行设置 nn.Softmax(dim-1) ) loaded_softmax_model.load_state_dict(torch.load(model_with_softmax.pth)) loaded_softmax_model.eval() # 推理示例 with torch.no_grad(): dummy_input torch.randn(1, *input_shape) # 替换为实际输入形状 probabilities loaded_softmax_model(dummy_input) print(probabilities)2、创建一个添加Softmax函数的包装模型import torch import torch.nn as nn import torch.nn.functional as F class MyModelWithSoftmax(nn.Module): def __init__(self, original_model_class_obj): super().__init__() # 实例化原始模型 self.backbone original_model_class_obj def forward(self, x): # 获取原始 logits logits self.backbone(x) # 应用 Softmax # 注意dim值需根据具体的模型识别场景进行设置 probs F.softmax(logits, dim-1) return probs # 1. 加载原始权重到 backbone original_model MyModel() # 假设这是你的原始模型类 original_model.load_state_dict(torch.load(model_weights.pth)) # 2. 创建新模型并复制权重 new_model MyModelWithSoftmax(original_model) # new_model.backbone.load_state_dict(original_model.state_dict()) new_model.eval() # 3. 保存新模型 torch.save(new_model.state_dict(), model_with_softmax_v2.pth)3、推理时动态添加import torch import torch.nn.functional as F # 1. 加载原始模型 model MyModel() model.load_state_dict(torch.load(model_weights.pth, map_locationcpu)) model.eval() # 2. 推理时动态转换 input_data torch.randn(1, 3, 224, 224) # 示例输入 with torch.no_grad(): logits model(input_data) # 输出原始分数 # 注意dim值需根据具体的模型识别场景进行设置 probabilities F.softmax(logits, dim-1) # 转换为概率 print(probabilities)补充知识Softmax函数在使用时需指定dim参数用于确定Softmax函数在输入向量的那个维度上做Softmax处理上文代码示例中Softmax函数dim参数设置使用时需结合具体情况指定。若经过Softmax后输出向量值均为1说明指定的维度dim上只有一个元素该情况下dim值设置的一定不正确可打印具体的输入向量及其形状shape结合模型的具体情况进行dim值的设置正确的dim设置输出向量的值小于1且在dim维度上的和为1若输出向量为多个维度上元素数量均小于1如维度为2x3输出向量在某个维度上的和为1dim的值设置也不一定正确需具体问题具体分析此处不做代码示例。代码示例如下In [1]: import torch In [2]: import torch.nn.functional as F In [3]: T_1x2 torch.Tensor([[3.4, 5.2]]) In [4]: print(T_1x2, T_1x2.shape) tensor([[3.4000, 5.2000]]) torch.Size([1, 2]) In [5]: # 错误dim设置 In [6]: F.softmax(T_1x2, dim0) Out[6]: tensor([[1., 1.]]) In [7]: # 正确dim设置 In [8]: F.softmax(T_1x2, dim1) Out[9]: tensor([[0.1419, 0.8581]])