欢迎您访问 最编程 本站为您分享编程语言代码,编程技术文章!
您现在的位置是: 首页

基于深度学习的毕尔巴鄂工程斑马鱼卵识别系统的注意事项和遇到的问题

最编程 2024-05-06 07:05:06
...

我是在云GPU上训练的,训练完成之后,保存了最优模型字典pth文件,在到cpu本地加载的时候,我们需要读取pth文件,正确代码如下

     # 加载模型权重,确保映射到正确的设备
                model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)  # 创建模型实例
                model=  torch.load('best-0.990.pth', map_location=device)

                model.to(device)  # 将模型映射到指定的设备

错误代码如下

import torchvision.models as models

# 创建模型架构的实例
model = models.resnet18(pretrained=False)  # 或者其他您使用的模型架构

# 加载模型的状态字典
model.load_state_dict(torch.load('best-0.990.pth', map_location=device), strict=False)

教程如下pytorch 状态字典:state_dict使用详解https://www.jb51.net/article/178720.htm