如何将经过 gpu 训练的模型加载到 cpu 中?
Posted
技术标签:
【中文标题】如何将经过 gpu 训练的模型加载到 cpu 中?【英文标题】:how to load the gpu trained model into the cpu? 【发布时间】:2019-08-25 23:14:39 【问题描述】:我正在使用 PyTorch。我将在带有 CPU 的多个 GPU 上使用已经训练好的模型。怎么做这个任务?
我在 Anaconda 3 和带有 cpu 的 pytorch 上试过,但我没有 gpu
model = models.get_pose_net(config, is_train=False)
gpus = [int(i) for i in config.GPUS.split(',')]
model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
print('Created model...')
print(model)
checkpoint = torch.load(config.MODEL.RESUME)
model.load_state_dict(checkpoint)
model.eval()
print('Loaded pretrained weights...')
我得到的错误是
AssertionError Traceback (most recent call last)
<ipython-input-15-bbfcd201d332> in <module>()
2 model = models.get_pose_net(config, is_train=False)
3 gpus = [int(i) for i in config.GPUS.split(',')]
----> 4 model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
5 print('Created model...')
6 print(model)
C:\Users\psl\Anaconda3\lib\site-packages\torch\nn\modules\module.py in cuda(self, device)
258 Module: self
259 """
--> 260 return self._apply(lambda t: t.cuda(device))
261
262 def cpu(self):
C:\Users\psl\Anaconda3\lib\site-packages\torch\nn\modules\module.py in
_apply(self, fn) 第185章 186 对于 self.children() 中的模块: --> 187 模块._apply(fn) 188 self._parameters.values() 中的参数为 189:
C:\Users\psl\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)
185 def _apply(self, fn):
186 for module in self.children():
--> 187 module._apply(fn)
188
189 for param in self._parameters.values():
C:\Users\psl\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)
191 # Tensors stored in modules are graph leaves, and we don't
192 # want to create copy nodes, so we have to unpack the data.
--> 193 param.data = fn(param.data)
194 if param._grad is not None:
195 param._grad.data = fn(param._grad.data)
C:\Users\psl\Anaconda3\lib\site-packages\torch\nn\modules\module.py in <lambda>(t)
258 Module: self
259 """
--> 260 return self._apply(lambda t: t.cuda(device))
261
262 def cpu(self):
C:\Users\psl\Anaconda3\lib\site-packages\torch\cuda\__init__.py in _lazy_init()
159 raise RuntimeError(
160 "Cannot re-initialize CUDA in forked subprocess. " + msg)
--> 161 _check_driver()
162 torch._C._cuda_init()
163 _cudart = _load_cudart()
C:\Users\psl\Anaconda3\lib\site-packages\torch\cuda\__init__.py in _check_driver()
80 Found no NVIDIA driver on your system. Please check that you
81 have an NVIDIA GPU and installed a driver from
---> 82 http://www.nvidia.com/Download/index.aspx""")
83 else:
84 # TODO: directly link to the alternative bin that needs install
AssertionError:
Found no NVIDIA driver on your system. Please check that you
have an NVIDIA GPU and installed a driver from
http://www.nvidia.com/Download/index.aspx
【问题讨论】:
【参考方案1】:要将保存的模型强制加载到 cpu 上,请使用以下命令。
torch.load('/path/to/saved/model', map_location='cpu')
在您的情况下将其更改为
torch.load(config.MODEL.RESUME, map_location='cpu')
【讨论】:
以上是关于如何将经过 gpu 训练的模型加载到 cpu 中?的主要内容,如果未能解决你的问题,请参考以下文章
[深度学习] Pytorch—— 多/单GPUCPU,训练保存加载模型参数问题
PyTorch笔记: GPU上训练的模型加载到CPU/错误处理Attempting to deserialize object on a CUDA device but torch.cuda.is_a