pytorch--load()模型参数加载

加载模型,一部分用于测试阶段,主要知识点在于数据加载的问题【torch.device】
官网链接1
官网链接2

不解释save
代码示例:

>>> torch.load('tensors.pt')
# Load all tensors onto the CPU
>>> torch.load('tensors.pt', map_location=torch.device('cpu'))
# Load all tensors onto the CPU, using a function
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Load all tensors onto GPU 1
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
# Load tensor from io.BytesIO object
>>> with open('tensor.pt', 'rb') as f:
        buffer = io.BytesIO(f.read())
>>> torch.load(buffer)

这里需要解释的是,不论之前save时tensor是cpu还是gpu,load时第一步总是在cpu上进行反序列化,如果之前是gpu那么load函数内部自动再继续映射到gpu上。
问题来了,如果你在github上找到的一些开源项目是gpu,而你的设备是cpu那么就会报错【因为你没有gpu】,这个时候便需要 map_location
下面是具体的示例


            if USE_CUDA:
                print("MODEL {} LOADED".format(str(path)))
                # 自动加载至cuda,不需要使用map_location
                self.encoder = torch.load(str(path)+'/enc.th')
                self.extKnow = torch.load(str(path)+'/enc_kb.th')
                self.decoder = torch.load(str(path)+'/dec.th')
            else:
                print("MODEL {} LOADED".format(str(path)))
                # 么得gpu,便需要声明使用cpu
                self.encoder = torch.load(str(path)+'/enc.th',lambda storage, loc: storage)
                self.extKnow = torch.load(str(path)+'/enc_kb.th',lambda storage, loc: storage)
                self.decoder = torch.load(str(path)+'/dec.th',lambda storage, loc: storage)

来源:【2019-ICLR: 任务型多轮对话系统】https://github.com/jasonwu0731/GLMP/blob/master/models/GLMP.py

发布了48 篇原创文章 · 获赞 9 · 访问量 1万+

猜你喜欢

转载自blog.csdn.net/NewDreamstyle/article/details/99003285
今日推荐