將TF模型權重存儲在CPU中嗎？

Question

我正在嘗試將為MNIST構建的CNN改編為CIFAR10。 我修改了結構以適應CIFAR圖像的大小（32x32）。 該模型在我的GPU上訓練，但是每當我運行測試數據的評估函數（ test_eval ）時，第一次卷積（ hconv1 ）就會出現內存不足錯誤-

 ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[10000,32,32,32]
     [[Node: Conv2D = Conv2D[T=DT_FLOAT, data_format="NHWC", padding="SAME", strides=[1, 1, 1, 1], use_cudnn_on_gpu=true, _device="/job:localhost/replica:0/task:0/gpu:0"](Reshape, W_conv1/read/_79)]]

現在，我嘗試在CPU上運行相同的代碼（通過使用export CUDA_VISIBLE_DEVICES="" ），並且test_eval函數成功運行。 所以我想，如果我僅將權重變量明確地存儲在CPU上，那可能會起作用。 所以我從-修改了weight_variable函數

def weight_variable(self, shape):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial)

至

def weight_variable(self, shape):
    with tf.device('/cpu:0'):
        initial = tf.Variable(tf.truncated_normal(shape, stddev=0.1))            
    return initial

但是，當我運行test_eval函數時，仍然出現OOM錯誤。 上面的代碼實際上將變量存儲在CPU中嗎？ 我該如何解決？

我的GPU是GeForce GTX 950M，2GB。 此外，對於MNIST CNN我上面貼的代碼是從適應這種

Answer 1

內存通常由激活占用。 例如，您的第一個轉換層的權重為3.2K，但激活次數為32x28x28x10000x4 = 1.003GB。 接下來是廣播添加，因此您需要超過2GB的RAM才能通過第一層。 請參閱此作為分析內存的示例