TensorFlow.Keras机器学习项目内存溢出问题排查求助

在TensorFlow.Keras机器学习项目中，处理完首批150张训练图像批次后，模型启动下一批次处理，打印1-3张调试图像后，RAM占用从7GB骤升至36GB，程序崩溃，始终无法进入验证阶段。

错误日志 #

Allocation of 5910240000 exceeds 10% of free system memory. This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. You must feed a value for placeholder tensor 'Placeholder/_0' with dtype int32. TF-TRT Warning: Could not find TensorRT

自定义data_generator代码 #

!pip install opencv-python
!pip install pydicom

import numpy as np
import pydicom
import cv2


def conditional_resize(img, img_id, target_height=3518, target_width=2800):
    # Get the original image dimensions
    original_height, original_width = img.shape[:2]

    # Initialize the output image as the input image
    output_img = img

    # Case 1: If the original image matches the target dimensions, leave it as it is
    if original_height == target_height and original_width == target_width:
        output_img = img

    else:
        # Case 2: If the original image is smaller than the target dimensions, pad it
        if original_height < target_height and original_width < target_width:
            pad_h = target_height - original_height
            pad_w = target_width - original_width
            output_img = cv2.copyMakeBorder(img, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0)

        # Case 3: If the original image is larger than the target dimensions, crop it
        elif original_height > target_height and original_width > target_width:
            crop_h = (original_height - target_height) // 2
            crop_w = (original_width - target_width) // 2
            output_img = img[crop_h:crop_h + target_height, crop_w:crop_w + target_width]

        # Case 4: If original image HEIGHT is larger than target AND WIDTH is smaller than target
        elif original_height > target_height and original_width < target_width:
            # Crop the height
            crop_h = (original_height - target_height) // 2
            cropped_img = img[crop_h:crop_h + target_height, :]

            # Pad the width
            pad_w = target_width - original_width
            output_img = cv2.copyMakeBorder(cropped_img, 0, 0, 0, pad_w, cv2.BORDER_CONSTANT, value=0)

        # Case 5: If original image HEIGHT is smaller than target AND WIDTH is larger than target
        elif original_height < target_height and original_width > target_width:
            # Pad the height
            pad_h = target_height - original_height
            padded_img = cv2.copyMakeBorder(img, 0, pad_h, 0, 0, cv2.BORDER_CONSTANT, value=0)

            # Crop the width
            crop_w = (original_width - target_width) // 2
            output_img = padded_img[:, crop_w:crop_w + target_width]

    print("\n IMAGE:", img_id, original_height, original_width, output_img.shape)
    return output_img

# 原代码中此部分被错误嵌套在conditional_resize函数内部（return语句之后），永远不会执行
mapping_dict_density = {'A': 0, 'B': 1, 'C': 2, 'D': 3}
mapping_dict_laterality = {'left': 0, 'right': 1}
mapping_dict_view_position = {'CC': 0, 'MLO': 1}

def my_data_generator(df, batch_size, sample_weights_birads, sample_weights_density):
    while True:
        # Select files (IDs) and labels for the batch
        batch_indices = np.random.choice(a=len(df), size=batch_size)
        batch_df = df.iloc[batch_indices]

        height = 3518  # The height of your images
        width = 2800  # The width of your images
        num_features = 2  # Number of feature columns, here you have 'laterality' and 'view_position'

        # Initialize your arrays
        batch_images = np.zeros((batch_size, height, width), dtype=np.float32)  # Adjust dtype as needed
        batch_labels_birads = np.zeros((batch_size, 1), dtype=np.int)  # Assuming birads labels are integers
        batch_labels_density = np.zeros((batch_size, 1), dtype=np.int)  # Assuming density labels are integers
        batch_features = np.zeros((batch_size, num_features), dtype=np.int)  # Assuming features are integers
        batch_weights_birads = np.zeros((batch_size,), dtype=np.float32)  # Assuming weights are float numbers
        batch_weights_density = np.zeros((batch_size,), dtype=np.float32)  # Assuming weights are float numbers


        for i, original_idx in enumerate(batch_indices):
            row = batch_df.iloc[i]

            img_id = row['image_id']
            study_id = row['study_id']
            img_path = f"/content/drive/MyDrive/Colab/vindr-mammo-a-large-scale-benchmark-dataset-for-computer-aided-detection-and-diagnosis-in-full-field-digital-mammography-1.0.0/images/{study_id}/{img_id}.dicom"
            img = pydicom.dcmread(img_path).pixel_array


            img = conditional_resize(img, img_id)
            batch_images[i] = img
            batch_labels_birads[i, 0] = row['breast_birads'] - 1
            batch_labels_density[i, 0] = mapping_dict_density.get(row['breast_density'], -1)
            batch_features[i] = [mapping_dict_laterality.get(row['laterality'], -1),
                                 mapping_dict_view_position.get(row['view_position'], -1)]
            batch_weights_birads[i] = sample_weights_birads[original_idx]
            batch_weights_density[i] = sample_weights_density[original_idx]



        yield {'image_input': batch_images, 'feature_input': batch_features}, \
            {'birads_output': batch_labels_birads, 'density_output': batch_labels_density}, \
            {'birads_output': batch_weights_birads, 'density_output': batch_weights_density}

Model.fit代码 #

history = model.fit(
    my_data_generator(
        training_data,
        batch_size=150,
        sample_weights_birads=sample_weights_birads_train,
        sample_weights_density=sample_weights_density_train,
    ),
    # steps_per_epoch=len(training_data) // 150,
    epochs=2,
    validation_data=my_data_generator(
        validation_data,
        batch_size=32,
        sample_weights_birads=sample_weights_birads_val,
        sample_weights_density=sample_weights_density_val,
    ),
    # validation_steps=len(validation_data) // 32,
    callbacks=[checkpoint]
)

可能的原因及修复方案 #

核心问题 #

1. 生成器语法致命错误
   原代码中mapping_dict定义和while True循环被错误嵌套在conditional_resize函数的return语句之后，这部分代码永远不会被执行，导致生成器逻辑完全失效，实际没有按预期生成批次，反而触发TensorFlow的占位符错误和内存异常。

2. 未设置训练步数导致内存堆积
   注释掉steps_per_epoch和validation_steps后，Keras会无限等待生成器输出批次，生成器不断生成新的张量但旧张量未被及时释放，最终内存被占满。

3. 批次尺寸过大
   单张3518×2800的float32图像占用约38.7MB内存，150张批次总内存约5.8GB，加上标签、权重和模型参数，接近日志中提示的5.5GB分配量，若叠加多个未释放的批次，内存会迅速突破上限。

4. 生成器内参数覆盖
   原代码在while循环内重新定义batch_size=150，会覆盖传入的验证集batch_size=32，导致验证集也生成大批次，进一步加剧内存压力。

修复步骤 #

• 修正语法错误：将mapping_dict和生成器函数（my_data_generator）移到conditional_resize函数外部，确保生成器逻辑正常运行。
• 恢复步数参数：取消steps_per_epoch和validation_steps的注释，设置为对应数据集大小除以批次尺寸，让Keras明确每个epoch的终止条件。
• 降低批次尺寸：将训练集batch_size从150降至32或64，观察内存占用变化，再逐步调整。
• 移除循环内的batch_size定义：使用函数传入的batch_size参数，保证训练/验证集批次大小符合预期。
• 关闭调试打印：注释掉conditional_resize中的print语句，减少IO资源消耗。
• 优化图像数据类型：若业务允许，将图像dtype从float32改为uint8，可将单张图像内存占用降至约9.7MB，大幅降低批次内存需求。

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