当前位置:网站首页>pytorch训练好的模型在加载和保存过程中的问题
pytorch训练好的模型在加载和保存过程中的问题
2022-07-06 08:27:00 【MAR-Sky】
在gpu上训练完成,在cpu上加载
torch.save(model.state_dict(), PATH)# 在gpu上训练后保存
# 在cpu的模型上加载使用
model.load_state_dict(torch.load(PATH, map_location='cpu'))
在cpu上训练完成,在gpu上加载
torch.save(model.state_dict(), PATH)# 在gpu上训练后保存
# 在cpu的模型上加载使用
model.load_state_dict(torch.load(PATH, map_location='cuda:0'))
在使用中需要注意的加载内容
当数据放入GPU,需要训练的模型也要放入GPU
''' data_loader:pytorch中加载数据 '''
for i, sample in enumerate(data_loader): # 对数据进行按批次遍历
image, target = sample # 每一批次加载返回值
if CUDA:
image = image.cuda() # 输入输出传入gpu
target = target.cuda()
# print(target.size)
optimizer.zero_grad() # 优化函数
output = mymodel(image)
mymodel.to(torch.device("cuda"))
多个gpu训练时的加载
参考:https://blog.csdn.net/weixin_43794311/article/details/120940090
import torch.nn as nn
mymodel = nn.DataParallel(mymodel)
pytorch中的nn模块使用nn.DataParallel将模型加载到多个GPU,需要注意,这种加载方式保存的权重参数会比不使用nn.DataParallel加载模型保存的权重参数的关键字前多一个"module."。是否使用nn.DataParallel加载模型,会导致下次再加载模型的时候可能会出现下图的问题,
当权重参数前面多一个“module."时,最简单的方式就是使用nn.DataParallel对模型加载,
边栏推荐
- China vanadium battery Market Research and future prospects report (2022 Edition)
- JVM 快速入门
- 面向个性化需求的在线云数据库混合调优系统 | SIGMOD 2022入选论文解读
- On the inverse order problem of 01 knapsack problem in one-dimensional state
- [MySQL] lock
- 2022 Inner Mongolia latest water conservancy and hydropower construction safety officer simulation examination questions and answers
- Use Alibaba icon in uniapp
- 【MySQL】日志
- Ruffian Heng embedded bimonthly, issue 49
- China's high purity aluminum target market status and investment forecast report (2022 Edition)
猜你喜欢
2022.02.13 - NC001. Reverse linked list
【MySQL】数据库的存储过程与存储函数通关教程(完整版)
JS native implementation shuttle box
[MySQL] database stored procedure and storage function clearance tutorial (full version)
vulnhub hackme: 1
Configuring OSPF load sharing for Huawei devices
What is the use of entering the critical point? How to realize STM32 single chip microcomputer?
JVM performance tuning and practical basic theory - Part 1
[research materials] 2022 China yuancosmos white paper - Download attached
游戏解包的危害及资源加密的重要性
随机推荐
Precise query of tree tree
LDAP application (4) Jenkins access
Analysis of pointer and array written test questions
C language custom type: struct
【MySQL】鎖
Grayscale upgrade tidb operator
CISP-PTE实操练习讲解
【Nvidia开发板】常见问题集 (不定时更新)
2022.02.13 - NC004. Print number of loops
JVM performance tuning and practical basic theory - Part 1
从表中名称映射关系修改视频名称
面向个性化需求的在线云数据库混合调优系统 | SIGMOD 2022入选论文解读
China Light conveyor belt in-depth research and investment strategy report (2022 Edition)
化不掉的钟薛高,逃不出网红产品的生命周期
Research Report on Market Research and investment strategy of microcrystalline graphite materials in China (2022 Edition)
使用 Dumpling 备份 TiDB 集群数据到兼容 S3 的存储
ROS编译 调用第三方动态库(xxx.so)
Function coritization
使用 BR 恢复 S3 兼容存储上的备份数据
Pointer advanced --- pointer array, array pointer