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

LansinBlog

于 2024-08-29 22:12:23 发布

阅读量30

点赞数

分类专栏： pytorch 文章标签： pytorch 深度学习

本文链接：https://blog.csdn.net/weixin_45674669/article/details/141689651

版权

pytorch 专栏收录该内容

8 篇文章 0 订阅

订阅专栏

1.模型权重保存 torch.save

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

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)