使用 tf.keras 模型训练时 GPU 随机冻结

Question

正在使用的版本： tensorflow-gpu: 2.0, CUDA v10, CuDNN v7.6.5, Python 3.7.4

系统规格： i9-7920X、4 个 RTX 2080Ti、128GB 2400MHz 内存、2TB SATA SSD

问题：

在使用 tensorflow 2.0 训练任何模型时，随机在一个 epoch 期间，GPU 会冻结并且 GPU 的功率使用将下降到 70W 左右，核心使用率为 0，内存利用率也固定在某个随机值。发生这种情况时，我也没有收到任何错误或异常。唯一的恢复方法是重新启动 jupyter 内核并从头开始运行。我首先想到我的代码可能有问题。所以我想我会尝试在 Cifar100 上训练 Densenet 时重现这个问题，但问题仍然存在。

如果我运行在多个 GPU 上进行训练，那么 GPU 也会冻结，但这种情况很少发生。但是使用单个 GPU，肯定会在某个时候卡住。

下面是用于训练Cifar100的代码

from densenet import DenseNet
from tensorflow.keras.datasets import cifar100
import tensorflow as tf
from tqdm import tqdm_notebook as tqdm

# the data, shuffled and split between train and test sets
(X_train, y_train), (X_test, y_test) = cifar100.load_data(label_mode='fine')
num_classes = 100
y_test_original = y_test

# Convert class vectors to binary class matrices. [one hot encoding]
y_train = tf.keras.utils.to_categorical(y_train, num_classes)
y_test = tf.keras.utils.to_categorical(y_test, num_classes)

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')

for i in range(3):
    mean = np.mean(X_train[:,:,:,i])
    std = np.std(X_train[:,:,:,i])
    X_train[:,:,:,i] = (X_train[:,:,:,i] - mean)/std
    X_test[:,:,:,i] = (X_test[:,:,:,i] - mean)/std


with tf.device('/gpu:0'):  
    model = DenseNet(input_shape=(32,32,3), dense_blocks=3, dense_layers=-1, growth_rate=12, nb_classes=100, dropout_rate=0.2,
             bottleneck=True, compression=0.5, weight_decay=1e-4, depth=100)


optimizer = tf.keras.optimizers.SGD(learning_rate=0.01, 
                                    momentum=0.9, 
                                    nesterov=True,
                                    name='SGD')
model.compile(loss = 'categorical_crossentropy', optimizer = optimizer, metrics = ['accuracy'])

def scheduler(epoch):
    if epoch < 151:
        return 0.01
    elif epoch < 251:
        return 0.001
    elif epoch < 301:
        return 0.0001

callback = tf.keras.callbacks.LearningRateScheduler(scheduler)

model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=300, verbose = 2)

PS：我什至在我的笔记本电脑上尝试了代码，它有一个 i7-8750h 和一个带有 32GB 和 970 EVO NVME 的 RTX 2060。不幸的是，我遇到了同样的 GPU 冻结问题。

有人知道问题出在哪里吗？

Answer 1

可能与包版本有关，请尝试降级到tensorflow 1.14。

也有可能是你使用了gpu重定位。试试在 cpu 上试一下：

 with tf.device('/cpu:0'):
    model = DenseNet(....)

parallel_model = multi_gpu_model(model, gpus=4)

然后可以使用parallel_model进行训练等等

Answer 2

所以我想 post 对此进行更新。 GPU 同步和 GPU 冻结问题终于解决了。以下是我所做的事情。我不知道是哪一个修复了它，或者他们是否都对修复做出了贡献：

DDU GPU 驱动程序并重新安装库存 CUDA v10 驱动程序。
更新主板 BIOS
使用下面的代码。

configproto = tf.compat.v1.ConfigProto() 
configproto.gpu_options.allow_growth = True
sess = tf.compat.v1.Session(config=configproto) 
tf.compat.v1.keras.backend.set_session(sess)

使用 tf.keras 模型训练时 GPU 随机冻结

GPU freezes randomly while training using tf.keras models

tensorflow

keras

nvidia

freeze

tensorflow2.0