求助：Keras 3搭配PyTorch后端训练速度较TensorFlow慢10倍

核心优化方向与操作步骤 #

1. 对齐数据格式与设备，消除跨设备拷贝开销 #

当前使用NumPy数组作为输入，PyTorch后端会频繁在CPU和GPU间拷贝数据，这是最大性能瓶颈。将数据转为PyTorch张量并直接部署到GPU：

# 替换原数据转换代码
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x_train = torch.tensor(x_train, dtype=torch.float32).to(device)
y_train = torch.tensor(y_train, dtype=torch.float32).to(device)
x_val = torch.tensor(x_val, dtype=torch.float32).to(device)
y_val = torch.tensor(y_val, dtype=torch.float32).to(device)

2. 削减回调带来的额外开销 #

BackupAndRestore在Windows+PyTorch环境下会产生大量磁盘IO和模型序列化开销，先移除该回调并调整ModelCheckpoint的保存策略：

# 仅保留必要的模型保存回调
savecallback = ModelCheckpoint(basefolder+"/"+modelfile, save_best_only=True, monitor='val_loss', mode='min', verbose=1)
hist=model.fit(x_train, y_train, validation_data=(x_val, y_val), batch_size=batchsize, epochs=20, callbacks=[savecallback])

3. 启用PyTorch LSTM的原生CUDA优化 #

Keras 3对PyTorch LSTM的封装默认可能未开启CuDNN加速，手动指定实现参数：

# 替换原有LSTM层定义
reg=0.00001
model.add(keras.layers.LSTM(
    80, 
    return_sequences=True, 
    dropout=0.0, 
    kernel_regularizer=l2(reg), 
    recurrent_regularizer=l2(reg),
    input_shape=(x_train.shape[1], x_train.shape[2]),
    implementation=2,  # 启用PyTorch CuDNN优化实现
    use_bias=True
))
model.add(keras.layers.LSTM(
    80, 
    return_sequences=False, 
    dropout=0.0, 
    kernel_regularizer=l2(reg), 
    recurrent_regularizer=l2(reg),
    implementation=2
))

4. 用PyTorch DataLoader优化批量加载 #

如果GPU利用率偏低，说明数据加载是瓶颈，改用DataLoader实现高效批量数据处理：

from torch.utils.data import TensorDataset, DataLoader

train_dataset = TensorDataset(x_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=batchsize, shuffle=True, pin_memory=True)
val_dataset = TensorDataset(x_val, y_val)
val_loader = DataLoader(val_dataset, batch_size=batchsize, pin_memory=True)

# 基于DataLoader训练模型
hist = model.fit(train_loader, validation_data=val_loader, epochs=20, callbacks=[savecallback])

5. 验证PyTorch与CUDA版本兼容性 #

确保PyTorch 2.3.1安装的是适配Windows系统的对应CUDA版本，运行以下命令验证：

print(torch.version.cuda)
print(torch.backends.cudnn.version())

若版本不匹配，卸载后重新安装对应CUDA版本的PyTorch。

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