Optuna调优3DCNN语义分割时遭遇CUDA内存不足问题求助

问题 #

使用Optuna、MONAI和PyTorch对3DCNN语义分割模型做参数调优，前期部分试验正常，新增试验后持续出现CUDA内存不足错误。已尝试torch.cuda.empty_cache()清理缓存但无效，希望在不修改网络结构和参数范围的前提下完成调优。

代码实现 #

# Librerias para el manejo de archivos
import logging
import os
import sys
import tempfile
from glob import glob

#Librerias para operaciones generales
import nibabel as nib
import numpy as np
import torch
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import csv
from Utilities import * 
#from torch.utils.tensorboard import SummaryWriter

#MONAI
import monai
from monai.data import create_test_image_3d, list_data_collate, decollate_batch
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric
from monai.transforms import (
    Activations,
    EnsureChannelFirstd,
    AsDiscrete,
    Compose,
    LoadImaged,
    RandCropByPosNegLabeld,
    RandRotate90d,
    ScaleIntensityd,
    SaveImage,
)
from monai.visualize import plot_2d_or_3d_image
from monai.engines import get_devices_spec

#Optuna
import optuna
def define_3D_unet(trial):
    kernel = trial.suggest_categorical("kernel_size",[1,3,5])
    drop_out = trial.suggest_float("drop_out",0,0.6)
    #Unet
    MODEL = monai.networks.nets.UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=1,
    kernel_size=kernel,
    channels=(16,32,64,128,256),
    strides=(2,2,2,2),
    num_res_units=2,
    dropout = drop_out
    )

    return MODEL

def objective(trial):
    device ="cuda" if torch.cuda.is_available() else "cpu"
    torch.cuda.empty_cache()
    #Definimos el pipeline de postprocesamiento
    post_trans = Compose(
        [
            Activations(sigmoid=True),
            AsDiscrete(threshold=0.5)
        ]
    )
    model = define_3D_unet(trial).to(device)

    optimizer_name = trial.suggest_categorical("optimizer",["Adam","RMSprop","Adadelta","Adagrad"])
    batch_size = trial.suggest_categorical("batch_size",[1,2,4])

    optimizer = getattr(torch.optim,optimizer_name)(model.parameters())
    #Función de perdida
    loss_function = monai.losses.DiceLoss(sigmoid=True)
    epochs = 400

    #Definimos la métrica que vamos a utilizar
    dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
    train_loader, val_loader, test_loader = preparar_dataloaders(train_files=train_files,
                                                                            val_files=val_files,
                                                                            test_files=test_files,
                                                                            batch_size_train=batch_size,
                                                                            device=device)
    
        
    val_interval = 2

    train_loss_values = []
    #train_metric_values = list()
    val_metric_values = list()

    for i in range(epochs):
        
        torch.cuda.empty_cache()
        print("-" * 10)
        print(f"epoch {i + 1}/{epochs}")

        model.train()
        epoch_loss=0
        step=0
        for batch_data in train_loader:
            step += 1

            # Separamos los datos de entrenamiento en datos y etiquetas
            inputs, labels = batch_data["img"].to(device),batch_data["seg"].to(device)

            #Zero grad
            optimizer.zero_grad()

            # Forward Pass
            outputs = model(inputs)

            # Computamos la función de perdida
            loss = loss_function(outputs, labels)

            #Backward pass
            loss.backward()

            #Optimizer step
            optimizer.step()

            #Calculamos la perdida de cada epoca y la perdida promedio
            epoch_loss += loss.item()
            epoch_len = 16//train_loader.batch_size
            print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")

        epoch_loss/= step
        train_loss_values.append(epoch_loss)
        #train_metric_values.append(dice_metric.aggregate().item())
        print(f"epoch {i + 1} avarage loss: {epoch_loss:.4f}")
        print(f'Memory allocated: {torch.cuda.memory_allocated()}')
        print(f'Memory max allocated:{torch.cuda.max_memory_allocated()}')

        if (i + 1) % val_interval == 0:
            model.eval()
            with torch.no_grad():
                val_images = None
                val_labels = None
                val_outputs = None

                for val_data in val_loader:
                    val_images,val_labels = val_data["img"].to(device), val_data["seg"].to(device)
                    roi_size=(96,96,96)
                    sw_batch_size = 4

                    #Calculamos las predicciones para los datos de validación
                    val_outputs = sliding_window_inference(val_images,roi_size,sw_batch_size,model)
                    val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]

                    #Calculamos el valor de la métrica
                    dice_metric(y_pred=val_outputs, y=val_labels)

                #Guardamos el valor final de la metrica
                metric = dice_metric.aggregate().item()

                #Guardamos los valores obtenidos en el arreglo
                val_metric_values.append(metric)

                #Reiniciamos el estado de la métrica
                dice_metric.reset()

    dice_metric = DiceMetric(include_background=True,reduction="mean", get_not_nans=False)

    post_trans = Compose(
    [
        Activations(sigmoid=True),
        AsDiscrete(threshold=0.5)
    ]
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.eval()

    with torch.no_grad():
        for test_data in test_loader:
            test_images, test_labels = test_data["img"].to(device),test_data["seg"].to(device)

            roi_size=(96,96,96)

            sw_batch_size = 4

            test_preds = sliding_window_inference(test_images,roi_size,sw_batch_size,model)

            test_preds = [post_trans(j) for j in decollate_batch(test_preds)]

            test_labels = decollate_batch(test_labels)

            dice_metric(y_pred=test_preds,y=test_labels)
    
    trial.report(dice_metric.aggregate().item(),epochs)

    if trial.should_prune():
        raise optuna.exceptions.TrialPruned()
    
    return dice_metric.aggregate().item()

if __name__ == "__main__":
    #Imprimimos la configuración para comprobar que este todo bien
    monai.config.print_config()
    logging.basicConfig(stream=sys.stdout, level=logging.INFO)

    data_dir = "./files/data/volume"
    labels_dir = "./files/data/labels/edit"

    train_files, val_files, test_files = dividir_archivos(data_dir, labels_dir)
    #Definimos las keys para simplicidad
    keys = ["img","seg"]

    # Device: setup as cpu for the lab
    device ="cuda" if torch.cuda.is_available() else "cpu"
    print(f'Device: {device}')
    torch.cuda.empty_cache()
    print(f'Memory allocated: {torch.cuda.memory_allocated()}')
    print(f'Memory max allocated:{torch.cuda.max_memory_allocated()}')

    study = optuna.load_study(storage="sqlite:///db.sqlite3",
                                study_name="prepration_study_400")
    study.optimize(objective,n_trials = 29,show_progress_bar=True)

    pruned_trials=[t for t in study.trials if t.state==optuna.trial.TrialState.PRUNED]
    complete_trials=[t for t in study.trials if t.state==optuna.trial.TrialState.COMPLETE]

    print("Study statistics: ")
    print(f" Number of finished trials: {len(study.trials)}")
    print(f" Number of prunned trials: {len(pruned_trials)}")
    print(f" Number of complete trials: {len(complete_trials)}")

    print("Best trial:")
    trial = study.best_trial

    print(f" Value: {trial.value}")

    print(f" Params:")
    for key, value in trial.params.items():
        print(f"   {key}: {value}")

报错信息 #

[W 2024-05-15 20:29:42,837] Trial 39 failed with value None.                                                                                                                                             
  0%|                                                                                                                                                                              | 0/29 [00:18<?, ?it/s]
Traceback (most recent call last):
  File "files/Optuna_training.py", line 218, in <module>
    study.optimize(objective,n_trials = 29,show_progress_bar=True)
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/optuna/study/study.py", line 451, in optimize
    _optimize(
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/optuna/study/_optimize.py", line 62, in _optimize
    _optimize_sequential(
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/optuna/study/_optimize.py", line 159, in _optimize_sequential
    frozen_trial = _run_trial(study, func, catch)
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/optuna/study/_optimize.py", line 247, in _run_trial
    raise func_err
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/optuna/study/_optimize.py", line 196, in _run_trial
    value_or_values = func(trial)
  File "files/Optuna_training.py", line 116, in objective
    loss.backward()
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/torch/_tensor.py", line 513, in backward
    return handle_torch_function(
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/torch/overrides.py", line 1621, in handle_torch_function
    result = torch_func_method(public_api, types, args, kwargs)
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/monai/data/meta_tensor.py", line 282, in __torch_function__
    ret = super().__torch_function__(func, types, args, kwargs)
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/torch/_tensor.py", line 1418, in __torch_function__
    ret = func(*args, **kwargs)
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/torch/_tensor.py", line 522, in backward
    torch.autograd.backward(
  File "/home/notvenja24/anaconda3/envs/training/lib/python3.8/site-packages/torch/autograd/__init__.py", line 266, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 216.00 MiB. GPU 0 has a total capacity of 3.82 GiB of which 209.50 MiB is free. Including non-PyTorch memory, this process has 3.57 GiB memory in use. Of the allocated memory 2.34 GiB is allocated by PyTorch, and 295.47 MiB 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.

可行解决方案 #

1. 启用PyTorch内存扩展段配置 #

设置环境变量缓解内存碎片问题，可在运行脚本前执行：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

或在Python代码开头添加：

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

2. 精细化训练循环内的内存释放 #

• 在每个训练batch结束后，手动删除无用张量并立即清理缓存：

for batch_data in train_loader:
    step += 1
    inputs, labels = batch_data["img"].to(device),batch_data["seg"].to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_function(outputs, labels)
    loss.backward()
    optimizer.step()
    epoch_loss += loss.item()
    
    # 释放当前batch的张量
    del inputs, labels, outputs, loss
    torch.cuda.empty_cache()

• 验证阶段结束后同样清理验证相关张量：

if (i + 1) % val_interval == 0:
    model.eval()
    with torch.no_grad():
        # ... 原有验证代码 ...
        dice_metric.reset()
    # 释放验证相关张量
    del val_images, val_labels, val_outputs
    torch.cuda.empty_cache()

3. 降低滑动窗口推理批次大小 #

验证和测试阶段的sw_batch_size=4占用内存过高，尝试改为2：

sw_batch_size = 2

4. 及时释放单trial的模型资源 #

在每个trial结束前，手动删除模型、优化器等大内存对象，避免内存残留：

def objective(trial):
    # ... 原有代码 ...
    final_metric = dice_metric.aggregate().item()
    
    # 清理当前trial的模型资源
    del model, optimizer, loss_function
    torch.cuda.empty_cache()
    
    if trial.should_prune():
        raise optuna.exceptions.TrialPruned()
    
    return final_metric

5. 强制Optuna单线程运行 #

设置n_jobs=1避免多trial并行占用内存：

study.optimize(objective, n_trials=29, show_progress_bar=True, n_jobs=1)

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