【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)
相关推荐
誉鏐19 分钟前
从零开始设计Transformer模型(1/2)——剥离RNN,保留Attention
人工智能·深度学习·transformer
神经星星35 分钟前
无需预对齐即可消除批次效应,东京大学团队开发深度学习框架STAIG,揭示肿瘤微环境中的详细基因信息
人工智能·深度学习·机器学习
程序员Linc1 小时前
写给新人的深度学习扫盲贴:向量与矩阵
人工智能·深度学习·矩阵·向量
补三补四1 小时前
机器学习-聚类分析算法
人工智能·深度学习·算法·机器学习
誉鏐2 小时前
PyTorch复现逻辑回归
人工智能·pytorch·逻辑回归
荷包蛋蛋怪2 小时前
【北京化工大学】 神经网络与深度学习 实验6 MATAR图像分类
人工智能·深度学习·神经网络·opencv·机器学习·计算机视觉·分类
贤小二AI3 小时前
贤小二c#版Yolov5 yolov8 yolov10 yolov11自动标注工具 + 免python环境 GPU一键训练包
人工智能·深度学习·yolo
意.远3 小时前
在PyTorch中使用GPU加速:从基础操作到模型部署
人工智能·pytorch·python·深度学习
Uzuki9 小时前
AI可解释性 II | Saliency Maps-based 归因方法(Attribution)论文导读(持续更新)
深度学习·机器学习·可解释性
byxdaz11 小时前
PyTorch中Linear全连接层
pytorch