Question: What takes so much gpu memory?

username: nanachiy
cooked: <p>I am a college student majoring in astronomy, and would like to do FFT and de-dispersion in GPU with torch module. The original data is a time series of 4GB in float32 format, and the factor of de-dispersion is the same. But it takes more than 24GB which exceeds my gpu memory. I am confused about what takes so much memory and how I can improve it. Thank you!</p>
<pre><code class="lang-auto">def dedisperse(spectrum,fre=5e8,DM=0.0,fre0=1e9,device='cuda', log_memory=False):   # do dedispersion for the signal
    # Create frequency data
    deltat = torch.arange(spectrum.shape[0], device='cuda:0', dtype = torch.float64)
    deltat = deltat * fre / spectrum.shape[0] + fre0
    print('deltat shape',deltat.shape)
    print('spectrum.shape[0]', spectrum.shape[0])
    phase=4.148808e15*(deltat**(-2)-fre0**(-2))*DM*deltat    # culculate the phase caused by the time-delay
    print('phase shape',phase.shape)
    print('phase:',phase)
    # DM
    del deltat
    torch.cuda.empty_cache()
    gc.collect()
    print('get_gpu_memory_del_deltat:',get_gpu_memory())
    spectrum = spectrum * torch.exp(-2j * math.pi * phase)
    spectrum_size_gb = spectrum.element_size() * spectrum.numel() / (1024 ** 3)
    print('get_gpu_memory_dedispers:',get_gpu_memory(),spectrum_size_gb,'GB',spectrum.element_size(),spectrum.numel())
    del phase
    torch.cuda.empty_cache()
    gc.collect()
    print('get_gpu_memory_del_phase:',get_gpu_memory())
    print('dedisperse completed')
    return spectrum#ddm_data
</code></pre>

username: Hamza_Javaid
cooked: <aside class="quote no-group" data-username="nanachiy" data-post="1" data-topic="221290">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nanachiy/48/77588_2.png" class="avatar"> nanachiy:</div>
<blockquote>
<pre><code class="lang-auto">def dedisperse(spectrum,fre=5e8,DM=0.0,fre0=1e9,device='cuda', log_memory=False):   # do dedispersion for the signal
    # Create frequency data
    deltat = torch.arange(spectrum.shape[0], device='cuda:0', dtype = torch.float64)
    deltat = deltat * fre / spectrum.shape[0] + fre0
    print('deltat shape',deltat.shape)
    print('spectrum.shape[0]', spectrum.shape[0])
    phase=4.148808e15*(deltat**(-2)-fre0**(-2))*DM*deltat    # culculate the phase caused by the time-delay
    print('phase shape',phase.shape)
    print('phase:',phase)
    # DM
    del deltat
    torch.cuda.empty_cache()
    gc.collect()
    print('get_gpu_memory_del_deltat:',get_gpu_memory())
    spectrum = spectrum * torch.exp(-2j * math.pi * phase)
    spectrum_size_gb = spectrum.element_size() * spectrum.numel() / (1024 ** 3)
    print('get_gpu_memory_dedispers:',get_gpu_memory(),spectrum_size_gb,'GB',spectrum.element_size(),spectrum.numel())
    del phase
    torch.cuda.empty_cache()
    gc.collect()
    print('get_gpu_memory_del_phase:',get_gpu_memory())
    print('dedisperse completed')
    return spectrum#ddm_data
</code></pre>
</blockquote>
</aside>
<p>When you compute <code>spectrum * torch.exp(-2j * math.pi * phase)</code>, PyTorch converts your float32 data to complex64 (or complex128), effectively doubling the memory requirement</p>
<ul>
<li><strong>Float64 operations</strong>: Your <code>deltat</code> and <code>phase</code> are computed in float64, which uses twice the memory of float32</li>
<li><strong>Intermediate tensors</strong>: The expression creates multiple intermediate tensors that aren’t immediately freed</li>
</ul>
<p>Try this</p>
<pre><code class="lang-auto">def dedisperse(spectrum, fre=5e8, DM=0.0, fre0=1e9, device='cuda'):
    # Use float32 consistently
    deltat = torch.arange(spectrum.shape[0], device=device, dtype=torch.float32)
    deltat = deltat * (fre / spectrum.shape[0]) + fre0
    
    # Calculate phase in float32
    phase = 4.148808e15 * (deltat**(-2) - fre0**(-2)) * DM * deltat
    
    # In-place operation to save memory
    spectrum = spectrum.to(torch.complex64)  # Convert once
    torch.exp(-2j * math.pi * phase, out=spectrum)  # Reuse memory
    
    return spectrum
</code></pre>

username: nanachiy
cooked: <p>Thank you for replying! In fact I made a mistake before: the program goes well when the data is in float32 format, but it warns me that the memory is not enough when it is float64. However, I calculate the total memory should be 16 GB while it takes more than 24 GB. I am confused what other thing takes so much memory.</p>