"RuntimeError: expected type torch.FloatTensor but got torch.cuda.FloatTensor" while resuming training

kirtac83: <p>Hi all,</p>
<p>I am trying to resume training from a pretrained Resnet-50 model, which is a 3dcnn model that has been initialized with Kinetics weights <a href="https://github.com/kenshohara/3D-ResNets-PyTorch" rel="nofollow noopener">https://github.com/kenshohara/3D-ResNets-PyTorch</a>. Initial training (5 epochs) has been done on a cuda device which has PyTorch version 1.0.0. My model and loss function is sent to current device, before starting training.<br>
I am unable to resume training on the same cuda machine, but I am able to resume it on a cpu device.<br>
The exact error output is as follows,</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/opt/anaconda3/lib/python3.7/site-packages/torch/optim/sgd.py", line 101, in step
    buf.mul_(momentum).add_(1 - dampening, d_p)
RuntimeError: expected type torch.FloatTensor but got torch.cuda.FloatTensor
</code></pre>
<p>I save states as follows,</p>
<pre><code class="lang-auto">        best_ckpt_path = os.path.join('checkpoint-best.tar')
        states = {
            'epoch': epoch + 1,
            'optimizer': optimizer.state_dict(),
            'state_dict': model.state_dict(),
        }
        torch.save(states, best_ckpt_path)
</code></pre>
<p>and I load,</p>
<pre><code class="lang-auto">        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        parameters = self.get_fine_tuning_parameters(opt.param_dict_list,
                                                     opt.learning_rate,
                                                     opt.weight_decay)

        if opt.optim.lower() in ['adam']:
            optimizer = optim.Adam(parameters, lr=opt.learning_rate, weight_decay=opt.weight_decay)
        elif opt.optim.lower() in ['sgd']:
            optimizer = optim.SGD(parameters, lr=opt.learning_rate, momentum=0.9, weight_decay=opt.weight_decay)
        else:
            raise ValueError('Invalid optimizer type string.')

        self.optimizer = optimizer

        if os.path.isfile(self.resume_path):
            print('resuming model from checkpoint {}'.format(self.resume_path))

            if self.device.type in ['cpu']:
                checkpoint = torch.load(self.resume_path, map_location=self.device)
            else:
                checkpoint = torch.load(self.resume_path)

            self.model.load_state_dict(checkpoint['state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer'])
            for state in self.optimizer.state.values():
                for k, v in state.items():
                    if isinstance(v, torch.Tensor):
                        print("device: {}".format(v.device))

            self.begin_epoch = checkpoint['epoch']
</code></pre>
<p>The print output of the above code block is <em>cpu</em> for each optimizer tensor, both for the cuda device and cpu device. So this shows that optimizer related tensors are already placed on cpu. So, how can I receive a <em>got torch.cuda.FloatTensor</em> error ?</p>

gfrogat: <p>Hi</p>
<p>Based on your error message, I guess that the data-types of the tensors in the <code>momentum_buffer</code> of the optimizer are on the wrong device.</p>
<aside class="quote no-group quote-modified" data-username="kirtac83" data-post="1" data-topic="37936">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kirtac83/48/12269_2.png" class="avatar"> kirtac83:</div>
<blockquote>
<p>Traceback (most recent call last): File “/opt/anaconda3/lib/python3.7/site-packages/torch/optim/sgd.py”, line 101, in step buf.mul_(momentum).add_(1 - dampening, d_p) RuntimeError: expected type torch.FloatTensor but got torch.cuda.FloatTensor</p>
</blockquote>
</aside>
<p>You can try to run this snippet to move the buffers:</p>
<pre data-code-wrap="python"><code class="lang-python">print(optimizer.state[list(optimizer.state.keys())[0]])

for p in optimizer.state.keys():
    param_state = optimizer.state[p]
    buf = param_state["momentum_buffer"]
    param_state["momentum_buffer"] = buf.cuda()  # move buf to device

print(optimizer.state[list(optimizer.state.keys())[0]])
</code></pre>
<p>See <a href="https://pytorch.org/docs/stable/_modules/torch/optim/sgd.html#SGD" rel="noopener nofollow ugc">here</a> how <code>momentum_buffers</code> are used in the <code>SGD</code> optimizer.</p>
<p>Hope this helps.</p>

kirtac83: <p>Hi,<br>
thanks a lot for the solution. I was not aware of the momentum buffers. And I haven’t specially treated them in my previous work. Why do you think this happened? To my knowledge, I only need to push the model and loss function to current device. Could you describe the best practice in this case please?</p>

gfrogat: <p>I’m not sure. If you call <code>torch.load</code> it should restore the tensors on the device prior to saving.<br>
<a href="https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer" rel="nofollow noopener">Optimizer</a> also casts the tensors to the respective device.</p>
<p>I can’t really tell without the complete code.</p>

Bananin: <p>I had a similar issue and I think it is related to your optimizer referencing model parameters that are not yet loaded on GPU. In my case I fixed the issue by changing the following:</p>
<pre><code class="lang-auto">optimizer = torch.optim.Adam(model.parameters(), lr=params.lr)
try:
       checkpoint = torch.load("checkpoints/"+model_folder+"checkpoint.pth.tar")
       start_epoch = checkpoint['epoch']
       scheduler.load_state_dict(checkpoint['scheduler'])
       model.load_state_dict(checkpoint['state_dict'])
       model.to(device) # moving this line here fixed the issue
       optimizer.load_state_dict(checkpoint['optimizer'])
       …
       # it used to be somewhere around here
</code></pre>
<p>So in your case a cleaner solution would be to send that <strong>parameters</strong> object to GPU (not sure how but there must be some one-liner that does that) BEFORE loading the dictionary, else it seems that it will be loaded to the usual memory instead of GPU memory.</p>