pytorch加载预训练网络权重,单GPU加载多GPU,CPU加载GPU权重

保存网络权重:

torch.save(G_AB.state_dict(), 'G_AB.pth')

单GPU加载权重:

G_AB.load_state_dict(torch.load(G_AB.pth))

CPU加载单GPU权重:

g.load_state_dict(torch.load('./save/G.pth',map_location='cpu')) #load weight

如果训练是多GPU,测试时用单GPU加载权重:

from collections import OrderedDict

state_dictBA = torch.load(opt.pth)
# create new OrderedDict that does not contain `module.`
new_state_dictBA = OrderedDict()
for k, v in state_dictBA.items():
    name = k[7:] # remove `module.`
    new_state_dictBA[name] = v
G_BA.load_state_dict(new_state_dictBA)

参考:https://www.ptorch.com/news/74.html

你可能感兴趣的:(其他)