0x00 前言

Pytorch裡使用optimizer的時候，由於其會記錄step等資訊， 有時會希望將optimizer的內容記錄下來，以備之後繼續使用， 那麼自然而然的會想到使用API中自帶的torch.save(object, path)torch.load(path)

再配合上optimizer.state_dict()optimizer.load_state_dict(obj) 來實現這一需求了~

於是，大家自然而然地會自信滿滿敲出如下這樣的語句——

torch.save(optimizer.state_dict(), path)
optimizer.load_state_dict(torch.load(path))
 

並收穫如下的Error——

RuntimeError                              Traceback (most recent call last)
<ipython-input-160-19f8d61b5e53> in <module>()
     37         optimizer.zero_grad()
     38         loss.backward()
---> 39         optimizer.step()
     40         print(model.state_dict()['linear_layer.weight'])
     41 

/usr/local/anaconda2/lib/python2.7/site-packages/torch/optim/adam.pyc in step(self, closure)
     63 
     64                 # Decay the first and second moment running average coefficient
---> 65                 exp_avg.mul_(beta1).add_(1 - beta1, grad)
     66                 exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
     67 

RuntimeError: arguments are located on different GPUs at /opt/conda/conda-bld/pytorch_1503966894950/work/torch/lib/THC/generated/../generic/THCTensorMathPointwise.cu:215
 

0x01 解決方案

二話不說上解決方案是我的習慣

# Load from dict
optimizer.load_state_dict(check_point['optim'])

# Load from file
optimizer.load_state_dict(torch.load(optim_path))

# Add this
for state in optimizer.state.values():
    for k, v in state.items():
        print (type(v))
        if torch.is_tensor(v):
            state[k] = v.cuda(cuda_id)
 

0x02 原理解釋

然後在慢慢的講為啥子~ 首先，這個方案是我在Issue中翻看到的： Thanks to pytorch/issues/2830

可以這麼理解，舉例說明，雖說你之前是放在GPU3上的，資料型別叫做 cuda.Tensor(GPU 3)， 但是天曉得你這個GPU3是哪臺機器上的GPU3哦，機器問了一下GPU3：是不是你家的啊， GPU3看了一眼計算完被打掃乾淨的戰場，已經空無一物——“不是吧，我家沒人啊”， 然後就委婉的拒絕了它。

所以，我們可以對load完畢的optimizer逐個詢問，只要是個tensor，我們就再把它介紹給GPU3一次~

--------------------- 本文來自 糖果天王 的CSDN 部落格 ，全文地址請點選：https://blog.csdn.net/okcd00/article/details/78470613?utm_source=copy