Torch框架下张量移至CPU后GPU内存无法释放的问题求助

Hey there, let's work through this GPU memory stuck issue you're hitting with Torch when moving tensors to CPU for large-scale neural style transfer. I’ve dealt with similar memory crunches in style transfer projects before, so here are some practical fixes to try:

1. Explicitly Nullify GPU Tensor References #

Just moving a tensor to the CPU doesn’t guarantee the original GPU tensor is released—Torch might still hold hidden references to it. You need to explicitly clear those references to let the garbage collector do its job:

if (torch.type(previousModule) == 'cudnn.ReLU' or torch.type(previousModule) == 'cudnn.SpatialConvolution') then
    -- Move output to CPU
    local cpu_output = previousModule.output:cpu()
    -- Replace the GPU tensor with the CPU version
    previousModule.output = cpu_output
    -- Nullify any potential hidden GPU references (check your module's properties!)
    if previousModule._gpu_output then
        previousModule._gpu_output = nil
    end
    -- Force garbage collection twice to ensure thorough cleanup
    collectgarbage('collect')
    collectgarbage('collect')
end

Many cudnn modules keep internal buffers for optimization—double-check your module’s attributes for any leftover GPU tensors you might need to nil out.

2. Use :clear() to Directly Free GPU Memory #

Torch’s CudaTensor has a clear() method that immediately releases the underlying GPU memory, which is more reliable than waiting for garbage collection. Try this workflow:

if (torch.type(previousModule) == 'cudnn.ReLU' or torch.type(previousModule) == 'cudnn.SpatialConvolution') then
    local gpu_tensor = previousModule.output
    -- Transfer to CPU first
    previousModule.output = gpu_tensor:cpu()
    -- Clear GPU memory right away
    gpu_tensor:clear()
    -- Remove all references to the GPU tensor
    gpu_tensor = nil
    -- Trigger garbage collection
    collectgarbage('collect')
end

This bypasses the garbage collector’s delay and frees memory immediately.

3. Clean Up Gradient Tensors in updateGradInput #

Don’t forget about gradient tensors! When working in updateGradInput, you need to apply the same cleanup logic to gradInput as you do to output:

if (torch.type(currentModule) == 'cudnn.ReLU' or torch.type(currentModule) == 'cudnn.SpatialConvolution') then
    if currentModule.gradInput then
        local gpu_grad = currentModule.gradInput
        -- Move gradient to CPU
        currentModule.gradInput = gpu_grad:cpu()
        -- Clear GPU gradient memory
        gpu_grad:clear()
        gpu_grad = nil
    end
    collectgarbage('collect')
end

Gradients can take up just as much memory as activations, so neglecting them will leave GPU memory tied up.

4. Verify Memory Usage to Debug #

To confirm your fixes are working, use cutorch to check GPU memory usage before and after transfers:

print("GPU Memory Before Transfer:", cutorch.getMemoryUsage(cutorch.getDevice()))
-- Run your transfer/cleanup code here
print("GPU Memory After Transfer:", cutorch.getMemoryUsage(cutorch.getDevice()))

This will help you pinpoint which layers or operations are failing to release memory, so you can adjust your cleanup logic accordingly.

内容的提问来源于stack exchange，提问作者Jareer