代码之家 › 专栏 › 技术社区 › Joseph Summerhays

如何缓存Pytorch模型,以便在未连接到互联网时使用?

vgg-net pytorch

0

Joseph Summerhays · 技术社区 · 6 年前

我在分类问题中使用vgg19。我可以使用校园研究计算机进行训练,但进行计算的节点无法访问互联网。所以运行一行代码,比如 self.net = models.vgg19(pretrained=True) 出现错误而失败 urllib.error.URLError: <urlopen error [Errno 101] Network is unreachable>

有没有一种方法可以将模型缓存在头节点(我可以访问互联网)上,并从缓存而不是计算节点上的互联网加载模型?

0 回复 | 直到 6 年前

1

4

unlut 6 年前

如果你只是将预训练网络的权重保存在某个地方,你就可以像加载任何其他网络权重一样加载它们。

保存:

import torchvision

#  I am assuming we have internet access here
model = torchvision.models.vgg16(pretrained=True)
torch.save(model.state_dict(), "Somewhere")

加载:

import torchvision

def create_vgg16(dict_path=None):
    model = torchvision.models.vgg16(pretrained=False)
    if (dict_path != None):
        model.load_state_dict(torch.load(dict_path))
    return model

model = create_vgg16("Somewhere")