如何按列拆分TensorFlow CNN模型并小批次训练以解决OOM错误？

看起来你遇到的核心问题是单样本特征维度过高（59536列），而不是样本数量多——常规的按行（样本）分批次根本没用，哪怕batch_size=1，单个样本的超大特征张量也会直接撑爆内存。下面给你一套具体的按特征列拆分模型+分批训练的方案：

核心思路 #

把近6万列的特征拆成若干个小特征块，让模型用多个独立的CNN分支分别处理每个特征块，最后再把所有分支的特征融合起来做分类/回归。这样每个批次里的张量都是小尺寸的，能有效降低内存占用。

具体代码修改 #

1. 先写特征拆分工具函数 #

这个函数会把高维特征按列拆分成多个子块，返回一个字典方便Estimator接收：

def split_high_dim_features(features, split_size=1024):
    # features shape: (样本数, 59536)
    total_cols = features.shape[1]
    num_splits = total_cols // split_size
    # 处理不能整除的情况，最后一块取剩余列
    if total_cols % split_size != 0:
        num_splits += 1
    # 按列拆分特征
    split_blocks = np.array_split(features, num_splits, axis=1)
    # 返回字典，键为x_0, x_1...x_n
    return {f"x_{idx}": block for idx, block in enumerate(split_blocks)}

你可以根据自己的GPU内存调整split_size，比如内存紧张就改成512甚至256。

2. 修改输入函数 #

不再传入整个超大特征张量，而是传入拆分后的特征块：

# 加载并拆分训练数据
train_x = np.array(training_set.data)
split_train_x = split_high_dim_features(train_x)
train_y = np.array(training_set.target)

train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x=split_train_x,
    y=train_y,
    num_epochs=None,
    batch_size=5,  # 现在按样本分批完全没问题，因为每个特征块都很小
    shuffle=True)

3. 重写CNN模型函数 #

让模型支持接收多个特征块，用分支处理后再融合：

def cnn_model_fn(features, labels, mode):
    # 提取所有拆分后的特征块
    feature_blocks = [features[f"x_{i}"] for i in range(len(features))]
    branch_outputs = []
    
    # 为每个特征块构建独立的CNN分支
    for block in feature_blocks:
        # 把一维特征reshape成适合1D卷积的形状：(批次大小, 特征块长度, 通道数)
        input_layer = tf.reshape(block, [-1, block.shape[1], 1])
        
        # 卷积+池化层（可以根据你的任务调整层数/参数）
        conv1 = tf.layers.conv1d(
            inputs=input_layer,
            filters=32,
            kernel_size=3,
            padding="same",
            activation=tf.nn.relu)
        pool1 = tf.layers.max_pooling1d(inputs=conv1, pool_size=2, strides=2)
        
        conv2 = tf.layers.conv1d(
            inputs=pool1,
            filters=64,
            kernel_size=3,
            padding="same",
            activation=tf.nn.relu)
        pool2 = tf.layers.max_pooling1d(inputs=conv2, pool_size=2, strides=2)
        
        # 扁平化特征，加入分支输出列表
        flat_feature = tf.layers.flatten(pool2)
        branch_outputs.append(flat_feature)
    
    # 融合所有分支的特征
    merged_features = tf.concat(branch_outputs, axis=1)
    
    # 后续全连接层和原来的逻辑一致
    dense = tf.layers.dense(inputs=merged_features, units=1024, activation=tf.nn.relu)
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
    
    # 输出层：替换成你的任务类别数
    logits = tf.layers.dense(inputs=dropout, units=YOUR_CLASS_NUM)
    
    # 以下是Estimator标准的预测/损失/训练逻辑，和你原来的代码兼容
    predictions = {
        "classes": tf.argmax(input=logits, axis=1),
        "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
    }
    
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
    
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
    
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    
    eval_metric_ops = {
        "accuracy": tf.metrics.accuracy(
            labels=labels, predictions=predictions["classes"])}
    return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

4. 进阶：用tf.data进一步优化内存 #

如果你的原始数据集加载到内存都困难，可以用tf.data.Dataset动态拆分特征，避免一次性加载整个大数组：

def split_features_tf(features):
    split_size = 1024
    total_cols = features.shape[1]
    num_splits = total_cols // split_size
    if total_cols % split_size != 0:
        num_splits += 1
    splits = tf.split(features, num_splits, axis=1)
    return {f"x_{i}": split for i, split in enumerate(splits)}

# 构建动态输入管道
dataset = tf.data.Dataset.from_tensor_slices((train_x, train_y))
dataset = dataset.map(lambda x, y: (split_features_tf(x), y))
dataset = dataset.shuffle(buffer_size=200).repeat().batch(5)

def train_input_fn():
    return dataset.make_one_shot_iterator().get_next()

额外注意事项 #

• 因为你的样本量很小（只有200行），训练时要注意过拟合问题，可以适当提高dropout比例，或者给卷积/全连接层加L2正则化。
• 如果你的特征是二维结构（比如图像拉伸成的一维向量），可以把卷积层改成2D的，同时调整reshape的形状适配。

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