如何将NIfTI(.nii)格式3D MRI传入TensorFlow训练并解决内存溢出问题

你代码不管模型搭得多小都崩溃，本质是两个完全踩中内存雷区的写法，和模型本身规模没有关系：

• 维度定义完全错位：你传入的input_shape=(193, 256, 256) 是整个3D脑体积的三维尺寸，但Conv2D层的输入要求是(图像高, 图像宽, 通道数)，TensorFlow会自动把第三个维度当成通道数——相当于你让网络一次性处理193x256大小、256个输入通道的张量，后面接Flatten层时，仅这一层的展开参数量就达到百万级，哪怕batch size设为1也会直接占满内存。
• 全量数据预加载的错误用法：tf.data.Dataset.from_tensor_slices 会把你传入的所有numpy数组整体常驻在程序内存里，哪怕单个NIfTI文件只有几十MB，攒上几十个被试的数据，直接就会把系统内存占满，触发Jupyter进程被系统杀掉。

根据你要搭建的网络类型，选对应的训练路径即可：

路径1：用2D卷积网络训练（对硬件要求低，适合入门跑通流程） #

核心逻辑是不要把整个3D脑体积当成一个样本，而是沿扫描轴把3D体积拆成独立的2D切片，每个切片单独作为一个训练样本。

• 优化思路：彻底抛弃“先把所有NIfTI读成numpy数组存在内存里”的写法，用懒加载管线实现“训练轮到哪个样本，才去硬盘读哪个文件、做预处理”，内存里只存文件路径，不存影像内容。
• 参考可运行的代码实现：

import nibabel as nib
import tensorflow as tf
import numpy as np

# 内存里只存文件路径和对应标签，不读影像内容
file_list = [("./data/subject_1.nii", 0), ("./data/subject_2.nii", 1)]

def load_preprocess(file_path, label):
    # 训练触发到这个文件时才读硬盘，不占常驻内存
    nii_array = nib.load(file_path.numpy().decode()).get_fdata()
    slice_samples = []
    slice_labels = []
    # 沿z轴拆分所有2D切片
    for idx in range(nii_array.shape[-1]):
        slice_img = nii_array[:, :, idx]
        # 灰度值归一化，补通道维变成单通道2D样本
        slice_img = (slice_img - slice_img.min()) / (slice_img.max() - slice_img.min() + 1e-8)
        slice_samples.append(slice_img[..., np.newaxis])
        slice_labels.append(label)
    return np.array(slice_samples), np.array(slice_labels)

# 构建懒加载数据集
dataset = tf.data.Dataset.from_tensor_slices(file_list)
dataset = dataset.map(
    lambda path, label: tf.py_function(
        func=load_preprocess,
        inp=[path, label],
        Tout=[tf.float32, tf.int64]
    ),
    num_parallel_calls=tf.data.AUTOTUNE
)
# 把每个文件拆出来的2D切片展平成独立训练样本
dataset = dataset.flat_map(
    lambda slices, labels: tf.data.Dataset.from_tensor_slices((slices, labels))
)
# 加batch、预取减少数据等待时间
dataset = dataset.batch(8).prefetch(tf.data.AUTOTUNE)

# 配套的小模型，输入维度适配单张2D切片
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(4, (3,3), padding='same', activation='relu',
                           input_shape=(193, 256, 1)),
    tf.keras.layers.MaxPooling2D((2, 2), strides=2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(2, activation='softmax')
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])
# 删掉原来固定step=1的错误写法，让网络遍历全部样本
model.fit(dataset, epochs=5)

路径2：用3D卷积网络训练（更适配3D空间特征，适合小样本研究） #

如果不想把3D影像拆成2D切片，就用3D卷积处理完整的脑体积，需要做三个内存优化操作：

• 输入维度必须补通道维：单个样本的正确shape是(193, 256, 256, 1)，不能把通道维弄丢。
• 对原始3D体积做降采样：把分辨率缩放到64x64x64或128x128x128，直接把单样本的数据量降低一个数量级。
• 同样用上面提到的懒加载数据集管线，不要一次性把所有影像读进内存；batch size设为1-2，配合梯度累积模拟更大batch的训练效果，避免内存溢出。

• 你原来代码里设置的steps_per_epoch=1 就算跑通了，也只会在单个batch上反复迭代，根本学不到有效特征，训练的时候删掉这个参数就行，让框架自动计算全量数据集的遍历步数。
• 不要试图靠缩小模型层的通道数解决内存问题——当前的内存溢出是数据加载和维度定义的问题，和模型大小无关，就算把Conv2D的通道数缩到1，错误的维度定义依然会导致内存爆炸。

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