TensorFlow不断消耗系统内存并在训练期间卡住

时间:2019-10-26 03:11:46

标签: tensorflow machine-learning deep-learning

我正在尝试训练一个小型网络以熟悉TensorFlow 2.0。但是,似乎tensorflow在我的计算机上无法正常工作。 这是我的代码:

import tensorflow as tf
from functools import reduce
from tensorflow.keras import layers, Sequential, datasets
import numpy as np




(train_images, train_labels), (test_images, test_labels) = datasets.cifar10.load_data()


train_images, test_images = train_images / 255.0, test_images / 255.0


class Model():
   weights = []
   biases = []

   def weights_collect(self):
       for l in self.layers:
           try:
               self.weights.append(l.kernel)
               self.biases.append(l.bias)
           except:
               pass

   def __init__(self):
       self.layers = [
           layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
           layers.MaxPool2D((2, 2)),
           layers.Conv2D(64, (3, 3), activation='relu'),
           layers.MaxPool2D((2, 2)),
           layers.Conv2D(64, (3, 3), activation='relu'),
           layers.MaxPool2D((2, 2)),
           layers.Flatten(),
           layers.Dense(1024, activation='relu'),
           layers.Dense(10)
       ]
       self.model = Sequential(self.layers)
       self.weights_collect()

   @tf.function
   def predict_logits(self, X):
       return self.model(X)


   @tf.function
   def __call__(self, X):
       return tf.nn.softmax(self.model(X))


@tf.function
def loss(m:Model, x:np.ndarray, t:np.ndarray):
   logits = m.predict_logits(x)
   tar = tf.one_hot(t, 10)
   return tf.math.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(tar, logits))

@tf.function
def acc(m, X, T):
   logits = m.predict_logits(X)
   target = tf.reshape(tf.dtypes.cast(T, tf.dtypes.int32), [-1])
   pred = tf.math.argmax(logits, axis=1, output_type=tf.dtypes.int32)
   return tf.math.reduce_sum(tf.dtypes.cast(pred==target, tf.dtypes.int32))/tf.shape(X)[0]


BATCH_SIZE = 1
dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).batch(BATCH_SIZE)


m = Model()
opt = tf.keras.optimizers.Adam()
EPOCHS = 20
for i in range(EPOCHS):
   for x, t in dataset:
       with tf.GradientTape() as tape:
           loss_value = loss(m, x, t)
       grads = tape.gradient(loss_value, m.model.trainable_weights)
       opt.apply_gradients(zip(grads, m.model.trainable_weights))
       print("hello")
   print(float(loss(m, test_images, test_labels)))
   print(float(acc(m, test_images, test_labels)))

运行此代码时,我不断收到这种错误消息:

Allocation of 1228800000 exceeds 10% of system memory.

然后我的模型将停止训练。

我试图更改批处理大小,但仍然无法正常工作。模型经过几次迭代训练后就死了。甚至我已将批处理大小更改为1。

TensorFlow似乎在训练期间一直在分配系统内存,而没有释放它。

我还重新安装了整个系统,试图解决此问题,但仍然无法正常工作。

4 个答案:

答案 0 :(得分:2)

刚刚解决了问题,最新版本的NVIDIA Game Ready Driver(440.97)似乎出了点问题。一旦回滚到436.48,即使上面提到的错误仍然存​​在,代码也将继续训练。

答案 1 :(得分:1)

您在计算机中使用哪个GPU?您的代码在其中装有GPU卡的情况下可以正常工作。我将批处理大小增加到128,并删除了print('Hello')语句,然后我就能够成功运行代码。

此外,由于您刚刚开始使用TensorFlow,GradientTape并不是开始学习TensorFlow的方法。它仅用于定制培训。利用model.fit()进行培训,或者从这些官方tutorials开始。

答案 2 :(得分:0)

我还使用Tensorflow对象检测api训练MaskRCNN模型时遇到了问题(第一步之后,它会挂起,没有使用GPU,并且CPU持续约18个)。该脚本在两周前运行良好。

长话短说,我回滚到建议的驱动程序,并且一切正常。我还尝试了最新更新(440.97之后的更新),问题仍然存在。

答案 3 :(得分:0)

我经历了同样的问题好几个星期了,最后

我明白了

我使用window10和Windows Server运行 使用powershell “。\ nvidia-smi -q -i 0 -d SUPPORTED_CLOCKS” 并看到NVidia驱动程序正在使用Cuda 10.2 当尝试将NVIDIA DRIVER降级为Cuda 10.1或CUDA 10.0时,它终于可以工作了 似乎受支持的CUDA10.2存在问题
enter image description here