【Pytorch】模型权重保存与上传

1.模型权重保存 torch.save

python 复制代码
model_name = args.model
if model_name == "ResNet18" or model_name == "ResNet34":
    from models.ResNet1 import BasicBlock
    from models.ResNet1 import ResNet as PATCHMODEL
if model_name == "ResNet18":
	net = PATCHMODEL(BasicBlock, [2, 2, 2, 2], num_classes=num_classes).cuda()
	
torch.save(net.state_dict(), weights_dir + '/' + model_name + '_train_loss_min_numCls{}.pth'.format(num_classes))

2.模型权重上传 load_state_dict

python 复制代码
model_name = args.model
if model_name == "ResNet18" or model_name == "ResNet34":
    from models.ResNet1 import BasicBlock
    from models.ResNet1 import ResNet as PATCHMODEL
if model_name == "ResNet18":
    model = PATCHMODEL(BasicBlock, [2, 2, 2, 2], num_classes=num_classes).cuda()
    
model.load_state_dict(torch.load(model_path), strict=False)
相关推荐
Luchang-Li3 分钟前
sglang pytorch NCCL hang分析
pytorch·python·nccl
Gyoku Mint7 小时前
提示词工程(Prompt Engineering)的崛起——为什么“会写Prompt”成了新技能?
人工智能·pytorch·深度学习·神经网络·语言模型·自然语言处理·nlp
m0_617663627 小时前
Deeplizard深度学习课程(七)—— 神经网络实验
人工智能·深度学习·神经网络
ningmengjing_8 小时前
激活函数:神经网络的“灵魂开关”
人工智能·深度学习·神经网络
THMAIL13 小时前
深度学习从入门到精通 - 迁移学习实战:用预训练模型解决小样本难题
人工智能·python·深度学习·算法·机器学习·迁移学习
却道天凉_好个秋15 小时前
计算机视觉(十二):人工智能、机器学习与深度学习
人工智能·深度学习·机器学习·计算机视觉
豆浩宇16 小时前
Conda环境隔离和PyCharm配置,完美同时运行PaddlePaddle和PyTorch
人工智能·pytorch·算法·计算机视觉·pycharm·conda·paddlepaddle
@国境以南,太阳以西16 小时前
基于Grad-CAM(Gradient-weighted Class Activation Mapping)的可解释性分析
人工智能·深度学习
格林威18 小时前
棱镜的技术加持:线扫相机如何同时拍RGB和SWIR?
人工智能·深度学习·数码相机·yolo·计算机视觉
Niuguangshuo18 小时前
深度学习:归一化技术
人工智能·深度学习