# Moving optimizer from CPU to GPU

**URL:** <https://discuss.pytorch.org/t/moving-optimizer-from-cpu-to-gpu/96068>\
**Category:** Uncategorized\
**Created:** [September 13, 2020, 12:01am UTC](https://discuss.pytorch.org/t/moving-optimizer-from-cpu-to-gpu/96068 "2020-09-13T00:01:05Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![amogkam](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/5f8ce5/32.png) [@amogkam](https://discuss.pytorch.org/u/amogkam)\
**Post date:** [September 13, 2020, 12:01am UTC](https://discuss.pytorch.org/t/moving-optimizer-from-cpu-to-gpu/96068/1 "2020-09-13T00:01:05Z")

</div>

I have a model and an optimizer and I want to save it’s state dict as CPU tensors. Then I want to load those state dicts back on GPU. This seems straightforward to do for a model, but what’s the best way to do this for the optimizer?

This is what my code looks like right now:

```auto
model = ...
optim = torch.optim.SGD(model.parameters(), momentum=0.1)

model_state = model.state_dict()
# Convert to CPU
for k, v in model_state.items():
   model_state[k] = v.cpu()

optim_state = optim.state_dict()
# Convert to CPU
for state in optim_state["state"].values():
    for k, v in state.items():
        state[k] = v.cpu()

# Now I want to load these state dicts back onto GPU
model2 = ...
model2.cuda()
optim2 = torch.optim.SGD(model2.parameters(), momentum=0.1)

# This seems to work; the model2 parameters are on GPU
model2.load_state_dict(model_state)

# Same does not hold true for optimizer
optim2.load_state_dict(optim_state)

```

The only option I see is to manually convert optimizer state back to Cuda

```auto
for state in optim2.state.values():
    for k, v in state.items():
        state[k] = v.cuda()

```

But would optim2 still update model2’s parameters?

---

<div class="post-metadata">

**Author:** ![aaniin](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aaniin/32/19065_2.png) [@aaniin](https://discuss.pytorch.org/u/aaniin)\
**Post date:** [September 16, 2020, 10:56am UTC](https://discuss.pytorch.org/t/moving-optimizer-from-cpu-to-gpu/96068/2 "2020-09-16T10:56:15Z")

</div>

I have the same problem.  
I have limited GPU memory. I can train with model and optimizer on GPU. However, GPU memory surges when loading model and optimizer to GPU, see [https://github.com/pytorch/pytorch/issues/7415](https://github.com/pytorch/pytorch/issues/7415)  
Effect is that I can’t load a previous checkpoint during training directly to GPU without going OOM. For the model, loading to CPU first and then moving to GPU works (see code below).  
Now, I go OOM when loading the optimizer. I would like to load the optimizer to CPU first, and then move it to the GPU. How can I do this?

```auto
# load model
# OOM triggered when directly loading to GPU, see https://github.com/pytorch/pytorch/issues/7415
# params = torch.load(model_save_path, map_location=lambda storage, loc: storage)
# Instead, load to CPU first
params = torch.load(model_save_path, map_location='cpu')
model.load_state_dict(params['state_dict'])
# And move model to GPU
model = model.to(device)

# Goes OOM - How can I load to CPU and move to GPU?
optimizer.load_state_dict(torch.load(model_save_path + '.optim'))

```

---

<div class="post-metadata">

**Author:** ![aaniin](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aaniin/32/19065_2.png) [@aaniin](https://discuss.pytorch.org/u/aaniin)\
**Post date:** [September 16, 2020, 11:38am UTC](https://discuss.pytorch.org/t/moving-optimizer-from-cpu-to-gpu/96068/3 "2020-09-16T11:38:28Z")

</div>

Hello @amogkam,  
Here [https://github.com/pytorch/pytorch/issues/8741](https://github.com/pytorch/pytorch/issues/8741) is an old feature request for a pytorch fct to move optimizer to device. I use the optimizer\_to function posted there and get around my OOM, training looks good so far.

```auto
def optimizer_to(optim, device):
    for param in optim.state.values():
        # Not sure there are any global tensors in the state dict
        if isinstance(param, torch.Tensor):
            param.data = param.data.to(device)
            if param._grad is not None:
                param._grad.data = param._grad.data.to(device)
        elif isinstance(param, dict):
            for subparam in param.values():
                if isinstance(subparam, torch.Tensor):
                    subparam.data = subparam.data.to(device)
                    if subparam._grad is not None:
                        subparam._grad.data = subparam._grad.data.to(device)

```

Here’s how I use it

```auto
# Load model
params = torch.load(model_save_path, map_location='cpu')
model.load_state_dict(params['state_dict'])
model = model.to(device)

# Empty any cache, not sure this helps, we try waht we can 
torch.cuda.empty_cache()

# Load optimizer
# Load to CPU first
optimizer.load_state_dict(torch.load(model_save_path + '.optim',map_location='cpu'))
# Send to GPU
optimizer_to(optimizer,device)

```

---

<div class="post-metadata">

**Author:** ![kirk86](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kirk86/32/151_2.png) [@kirk86](https://discuss.pytorch.org/u/kirk86)\
**Post date:** [May 9, 2021, 8:08pm UTC](https://discuss.pytorch.org/t/moving-optimizer-from-cpu-to-gpu/96068/4 "2021-05-09T20:08:03Z")

</div>

Best solution for this would be for pytorch to provide similar interface to `model.to(device)` for the optimizer `optim.to(device)` as well.

Another solution would have been to not save tensors in the state dicts with the device argument in them so that when loading a model would not result in this discrepancy between model state dict and optim state dict.

e.g. If all tensors in state dicts of both model and optim do not have a `device` assoctiated with them then simply doing

```auto
state_dict = torch.load('some_model.pt')
model.load_state_dict(statedict['model'])
optim.load_state_dict(state_dict['optim'])

```

would have worked fine .