如何在GPU训练的模型后用保存,然后用CPU版导入
收藏
如何在GPU训练的模型后用保存,然后用CPU版导入,像Pytorch一样,通过该设备就行?
0
收藏
请登录后评论
用***.state_dict()和***.set_state_dict()两个API。
例如:
from paddle import nn
from paddle import save,load
net=nn.Linear(3,3)
save(net.state_dict(),'lorem_ipsum.dolor') #先用net.state_dict()获取网络的参数列表,再将网络参数存储在文件'lorem_ipsum.dolor'中
sd=load('lorem_ipsum.dolor') #读取存储在文件'lorem_ipsum.dolor'中的参数列表
net.set_state_dict(sd) #将参数列表加载到网络中,这个API在pytorch中写作net.load_state_dict(sd)
和pytorch几乎一样的