【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)
相关推荐
孙同学要努力3 小时前
全连接神经网络案例——手写数字识别
人工智能·深度学习·神经网络
丕羽4 小时前
【Pytorch】基本语法
人工智能·pytorch·python
sniper_fandc5 小时前
深度学习基础—循环神经网络的梯度消失与解决
人工智能·rnn·深度学习
weixin_518285055 小时前
深度学习笔记10-多分类
人工智能·笔记·深度学习
阿_旭7 小时前
基于YOLO11/v10/v8/v5深度学习的维修工具检测识别系统设计与实现【python源码+Pyqt5界面+数据集+训练代码】
人工智能·python·深度学习·qt·ai
YRr YRr7 小时前
深度学习:Cross-attention详解
人工智能·深度学习
阿_旭7 小时前
基于YOLO11/v10/v8/v5深度学习的煤矿传送带异物检测系统设计与实现【python源码+Pyqt5界面+数据集+训练代码】
人工智能·python·深度学习·目标检测·yolo11
算家云7 小时前
如何在算家云搭建Aatrox-Bert-VITS2(音频生成)
人工智能·深度学习·aigc·模型搭建·音频生成·算家云
小言从不摸鱼8 小时前
【NLP自然语言处理】深入解析Encoder与Decoder模块:结构、作用与深度学习应用
人工智能·深度学习·神经网络·机器学习·自然语言处理·transformer·1024程序员节
湫ccc8 小时前
Bert框架详解(上)
人工智能·深度学习·bert