Map_location() takes multiple devices

username: tengerye
cooked: <p>I save a model from <code>cuda:0</code> (<code>model1</code>) and load it to <code>cuda:1</code> (<code>model2</code>). After I delete the tensor on <code>model1</code>, the <code>cuda:0</code> is still occupied. Why is that? The following is the snippet to follow.</p>
<pre><code class="lang-auto"># Define model
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        x = self.fc(x)
        return x

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

def torch_save(model, optimizer):
    torch.save({'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'epoch': 0
                }, 'test.pt')

def torch_load(optimizer):
    model = TheModelClass()

    checkpoint = torch.load('test.pt', map_location={'cuda:0': 'cuda:1'})
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']

    print(f'epoch = {epoch}')

    return model


model = TheModelClass()

model1 = model.to('cuda')
print(f'Device of model = {model.fc.weight.device}')

torch_save(model1, optimizer)
model2 = torch_load(optimizer)
del model, model1

torch.cuda.empty_cache()
# Expected to have only cuda:1 has tensors,
# but both cuda:0 and cuda:1 have tensors.
# This is wired.</code></pre>

username: ptrblck
cooked: <p>The first CUDA operation on a device will create the CUDA context (which contains the PyTorch kernels, cudnn, NCCL etc.). While you are able to delete all tensors, you are not able to release the CUDA context during the script lifetime.</p>

username: tengerye
cooked: <p>Hi, thank you for your kind reply. I thought <code>torch.cuda.empty_cache() </code> will release the context. But I found something even more wired: if I just load the <code>model2</code> (without loading <code>model1</code>). Two gpus are still used.</p>

username: ptrblck
cooked: <p><code>torch.cuda.empty_cache()</code> will only remove all reserved and unused memory, which the custom caching allocator in PyTorch holds to avoid expensive memory allocations. The CUDA context won’t be freed.</p>
<p>If I understand the <code>map_location</code> argument with a <code>dict</code>, it would remap the tensors from the <code>key</code> location to the <code>value</code> location, so both GPUs would be initially used (which should thus also create a context on both of them).<br>
To avoid it you could store the <code>state_dict</code> on the CPU and map it to the desired GPU afterwards.</p>

username: tengerye
cooked: <p>Hi, thank you for your help and patience. May I ask if it works in <code>DistributedDataParallel</code> as well? Since the <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">doc</a> reads</p>
<blockquote>
<p>Besides, when loading the module, you need to provide an appropriate  <code>map_location</code>  argument to prevent a process to step into others’ devices. If  <code>map_location</code>  is missing,  <code>torch.load</code>  will first load the module to CPU and then copy each parameter to where it was saved, which would result in all processes on the same machine using the same set of devices.</p>
</blockquote>

username: ptrblck
cooked: <p>You should be able to load a CPU <code>state_dict</code> to each device using DDP.<br>
However, if the original <code>state_dict</code> was stored on a specific GPU, you might run into the mentioned issue.<br>
Let me know, if this would be the case, as I haven’t seen it yet.</p>