torch.OutOfMemoryError: CUDA out of memory

Mohan_Krishna: <p><strong>torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 574.00 MiB. GPU 0 has a total capacity of 23.66 GiB of which 587.00 MiB is free. Including non-PyTorch memory, this process has 22.33 GiB memory in use. Of the allocated memory 17.67 GiB is allocated by PyTorch, and 4.39 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (<a href="https://pytorch.org/docs/stable/notes/cuda.html#environment-variables" class="inline-onebox" rel="noopener nofollow ugc">CUDA semantics — PyTorch 2.6 documentation</a>)</strong></p>
<p>Hi everyone, I am using NVIDIA RTX A5000 24GB memory. Currently I am using Driver Version: 550.120, cuda 12.1, cudnn=90100, pytorch=2.5.1. Can someone suggest : to fix the above issue</p>

ptrblck: <p>You could reduce the memory requirement e.g. by reducing the batch size or check if <code>expandable_segments</code> helps reducing the memory fragmentation as explained in the error message.</p>

Mohan_Krishna: <p>I have tried expandable_segmens=true, and trying with batch size =2. I am working on video data. Also I have already tried batch size=1. Any other suggestion please?</p>

Mohan_Krishna: <p>The following code i am running on NVIDIA RTX A5000 of 24 GB memory. If I use batch size=2, the cuda out of memory issue is coming.</p>
<p>import os<br>
import pdb<br>
import sys<br>
import copy<br>
import torch<br>
import numpy as np<br>
import torch.nn as nn<br>
from tqdm import tqdm<br>
import torch.nn.functional as F<br>
import matplotlib.pyplot as plt<br>
from evaluation.slr_eval.wer_calculation import evaluate<br>
from torch.cuda.amp import autocast as autocast<br>
from torch.cuda.amp import GradScaler</p>
<p>def seq_train(loader, model, optimizer, device, epoch_idx, recoder):<br>
model. Train()<br>
loss_value = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
clr = [group[‘lr’] for group in optimizer.optimizer.param_groups]<br>
scaler = GradScaler()<br>
for batch_idx, data in enumerate(tqdm(loader)):<br>
vid = device.data_to_device(data[0])<br>
vid_lgt = device.data_to_device(data[1])<br>
label = device.data_to_device(data[2])<br>
label_lgt = device.data_to_device(data[3])<br>
optimizer.zero_grad()<br>
with autocast():<br>
ret_dict = model(vid, vid_lgt, label=label, label_lgt=label_lgt)<br>
loss = model.criterion_calculation(ret_dict, label, label_lgt)<br>
if np.isinf(loss.item()) or np.isnan(loss.item()):<br>
print(‘loss is nan’)<br>
<span class="hashtag-raw">#print</span>(data[-1])<br>
print(str(data[1])+’  frames’)<br>
print(str(data[3])+’  glosses’)<br>
continue<br>
scaler.scale(loss).backward()<br>
scaler.step(optimizer.optimizer)<br>
scaler.update()<br>
# nn.utils.clip_grad_norm_(model.rnn.parameters(), 5)<br>
loss_value.append(loss.item())<br>
if batch_idx % recoder.log_interval == 0:<br>
recoder.print_log(<br>
‘\tEpoch: {}, Batch({}/{}) done. Loss: {:.8f}  lr:{:.6f}’<br>
.format(epoch_idx, batch_idx, len(loader), loss.item(), clr[0]))<br>
del ret_dict<br>
del loss<br>
optimizer.scheduler.step()<br>
recoder.print_log(‘\tMean training loss: {:.10f}.’.format(np.mean(loss_value)))<br>
return</p>
<p>def seq_eval(cfg, loader, model, device, mode, epoch, work_dir, recoder,<br>
evaluate_tool=“python”):<br>
model.eval()<br>
total_sent = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
total_info = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
total_conv_sent = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
stat = {i: [0, 0] for i in range(len(loader.dataset.dict))}<br>
for batch_idx, data in enumerate(tqdm(loader)):<br>
recoder.record_timer(“device”)<br>
vid = device.data_to_device(data[0])<br>
vid_lgt = device.data_to_device(data[1])<br>
label = device.data_to_device(data[2])<br>
label_lgt = device.data_to_device(data[3])<br>
with torch.no_grad():<br>
ret_dict = model(vid, vid_lgt, label=label, label_lgt=label_lgt)</p>
<pre><code>    total_info += [file_name.split("|")[0] for file_name in data[-1]]
    total_sent += ret_dict['recognized_sents']
    total_conv_sent += ret_dict['conv_sents']
try:
    python_eval = True if evaluate_tool == "python" else False
    write2file(work_dir + "output-hypothesis-{}.ctm".format(mode), total_info, total_sent)
    write2file(work_dir + "output-hypothesis-{}-conv.ctm".format(mode), total_info,
               total_conv_sent)
    conv_ret = evaluate(
        prefix=work_dir, mode=mode, output_file="output-hypothesis-{}-conv.ctm".format(mode),
        evaluate_dir=cfg.dataset_info['evaluation_dir'],
        evaluate_prefix=cfg.dataset_info['evaluation_prefix'],
        output_dir="epoch_{}_result/".format(epoch),
        python_evaluate=python_eval,
    )
    lstm_ret = evaluate(
        prefix=work_dir, mode=mode, output_file="output-hypothesis-{}.ctm".format(mode),
        evaluate_dir=cfg.dataset_info['evaluation_dir'],
        evaluate_prefix=cfg.dataset_info['evaluation_prefix'],
        output_dir="epoch_{}_result/".format(epoch),
        python_evaluate=python_eval,
        triplet=True,
    )
except:
    print("Unexpected error:", sys.exc_info()[0])
    lstm_ret = 100.0
finally:
    pass
del conv_ret
del total_sent
del total_info
del total_conv_sent
del vid
del vid_lgt
del label
del label_lgt
recoder.print_log(f"Epoch {epoch}, {mode} {lstm_ret: 2.2f}%", f"{work_dir}/{mode}.txt")
return lstm_ret
</code></pre>
<p>def seq_feature_generation(loader, model, device, mode, work_dir, recoder):<br>
model.eval()</p>
<pre><code>src_path = os.path.abspath(f"{work_dir}{mode}")
tgt_path = os.path.abspath(f"./features/{mode}")
if not os.path.exists("./features/"):
    os.makedirs("./features/")

if os.path.islink(tgt_path):
    curr_path = os.readlink(tgt_path)
    if work_dir[1:] in curr_path and os.path.isabs(curr_path):
        return
    else:
        os.unlink(tgt_path)
else:
    if os.path.exists(src_path) and len(loader.dataset) == len(os.listdir(src_path)):
        os.symlink(src_path, tgt_path)
        return

for batch_idx, data in tqdm(enumerate(loader)):
    recoder.record_timer("device")
    vid = device.data_to_device(data[0])
    vid_lgt = device.data_to_device(data[1])
    with torch.no_grad():
        ret_dict = model(vid, vid_lgt)
    if not os.path.exists(src_path):
        os.makedirs(src_path)
    start = 0
    for sample_idx in range(len(vid)):
        end = start + data[3][sample_idx]
        filename = f"{src_path}/{data[-1][sample_idx].split('|')[0]}_features.npy"
        save_file = {
            "label": data[2][start:end],
            "features": ret_dict['framewise_features'][sample_idx][:, :vid_lgt[sample_idx]].T.cpu().detach(),
        }
        np.save(filename, save_file)
        start = end
    assert end == len(data[2])
os.symlink(src_path, tgt_path)
</code></pre>
<p>def write2file(path, info, output):<br>
filereader = open(path, “w”)<br>
for sample_idx, sample in enumerate(output):<br>
for word_idx, word in enumerate(sample):<br>
filereader.writelines(<br>
“{} 1 {:.2f} {:.2f} {}\n”.format(info[sample_idx],<br>
word_idx * 1.0 / 100,<br>
(word_idx + 1) * 1.0 / 100,<br>
word[0]))</p>