训练Keras词性标注器时GPU内存分配错误的解决方法

训练Keras词性（POS）标注器时触发GPU内存分配错误，报错显示GPU尝试分配73.90GiB内存失败，最终提示“Failed copying input tensor from CPU to GPU”。

环境配置 #

• Python 3.9 + TensorFlow 2.10
• CUDA Toolkit 11.2 + cuDNN 8.2（Windows 10原生支持的最后一版配置）
• 硬件：NVIDIA GeForce RTX 2070 SUPER（8GB VRAM），主机内存64GB

数据与预处理 #

训练数据为词元和POS标签元组组成的句子列表：

[[("hello", "INTJ"), ("world", "NOUN"), ("!", "PUNCT")], [("oh", "INTJ"), ("hi", "INTJ")], ...]

数据划分训练、验证、测试集后，用sklearn的DictVectorizer完成向量化与独热编码。

模型代码 #

模型构建函数 #

def construct_model(input_dim, hidden_neurons, output_dim):
    pos_model = Sequential([
        Dense(hidden_neurons, input_dim=input_dim),
        Activation('relu'),
        Dropout(0.2),
        Dense(hidden_neurons),
        Activation('relu'),
        Dropout(0.2),
        Dense(output_dim, activation='softmax')
    ])
    pos_model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

    return pos_model

训练代码 #

if __name__ == "__main__":

    X_train = processed_data[0]
    y_train = processed_data[1]

    X_val = processed_data[2]
    y_val = processed_data[3]

    X_test = processed_data[4]
    y_test = processed_data[5]

    model_params = {
        'build_fn': construct_model,
        'input_dim': X_train.shape[1],
        'hidden_neurons': 512,
        'output_dim': y_train.shape[1],
        'epochs': 5,
        'batch_size': 256,
        'verbose': 1,
        'validation_data': (X_val, y_val),
        'shuffle': True
    }

    classifier = KerasClassifier(**model_params)
    pos_model = classifier.fit(X_train, y_train)

已尝试的无效方法 #

• 将batch size降至2
• 启用TensorFlow GPU内存增长设置
• 考虑使用gc.collect()但无循环场景无法应用

1. 压缩输入特征维度（核心解决方向） #

独热编码会因词汇量过大导致input_dim（特征数）爆炸，这是GPU内存溢出的主要原因：

• 过滤低频词：仅保留出现次数≥5的词汇，缩小词汇表规模
• 改用词嵌入替代独热编码：用Embedding层将高维独热向量映射到低维稠密向量（如100/200维），示例修改如下：

def construct_model(vocab_size, embedding_dim, hidden_neurons, output_dim):
    pos_model = Sequential([
        Embedding(input_dim=vocab_size, output_dim=embedding_dim, input_length=1),
        Flatten(),
        Dense(hidden_neurons, activation='relu'),
        Dropout(0.2),
        Dense(hidden_neurons, activation='relu'),
        Dropout(0.2),
        Dense(output_dim, activation='softmax')
    ])
    pos_model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    return pos_model

注意：使用Embedding层时，输入需为词的索引（而非独热向量），需用LabelEncoder对词汇做编码。

2. 强制限制GPU内存占用 #

直接指定TensorFlow可使用的GPU内存上限，避免抢占全部显存：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.set_logical_device_configuration(
            gpus[0],
            [tf.config.LogicalDeviceConfiguration(memory_limit=6144)]  # 限制为6GB，留2GB给系统
        )
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        print(e)

3. 简化模型结构 #

• 降低隐藏层神经元数量：从512降至256或128，减少模型参数总量
• 减少全连接层数量：去掉一组Dense+Dropout层，简化模型结构

4. 降低数据精度 #

将输入数据从float64转为float32，减少内存占用：

X_train = X_train.astype('float32')
y_train = y_train.astype('float32')
X_val = X_val.astype('float32')
y_val = y_val.astype('float32')

5. 分批加载数据 #

若数据集过大，改用生成器分批加载，避免一次性将所有数据放入内存：

from tensorflow.keras.utils import Sequence
import numpy as np

class POSDataGenerator(Sequence):
    def __init__(self, x_set, y_set, batch_size):
        self.x, self.y = x_set, y_set
        self.batch_size = batch_size

    def __len__(self):
        return int(np.ceil(len(self.x) / self.batch_size))

    def __getitem__(self, idx):
        batch_x = self.x[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_y = self.y[idx * self.batch_size:(idx + 1) * self.batch_size]
        return batch_x, batch_y

# 使用生成器训练
train_generator = POSDataGenerator(X_train, y_train, batch_size=32)
val_generator = POSDataGenerator(X_val, y_val, batch_size=32)
pos_model.fit(train_generator, validation_data=val_generator, epochs=5)

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