我在 TensorFlow 中有一个循环,如下所示:
with tf.device("/gpu:1"):
losses = []
for target, output in zip(targets, lstm_outputs):
logits = tf.matmul(W, output) + b
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits, target)
losses.append(loss)
total_loss = tf.add_n(losses)
为该层分配梯度时出现 OOM 错误,因为每个矩阵乘法都是占用内存的图中的不同操作。有没有办法阻止 TensorFlow 同时分配所有这些操作?
对于 TensorFlow 优化来说,这是一个具有挑战性的图,因为必须保留每一层的激活来聚合单个梯度W
。一种可能性是通过实验性的 https://github.com/tensorflow/tensorflow/blob/0249729a26b2cd3bdbbc5880f8d634a56860b0fd/tensorflow/python/ops/gradients.py#L580 aggregation_method
调用时的参数optimizer.optimize() https://www.tensorflow.org/versions/r0.7/api_docs/python/train.html#Optimizer.minimize.
例如,您可以尝试以下操作:
optimizer = tf.train.AdagradOptimizer(...) # Or another optimization algorithm.
train_op = optimizer.minimize(
total_loss,
aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)
此选项急切地聚合经常使用的变量的梯度,而不是将它们全部保留在内存中,直到计算出所有梯度。如果这不起作用,则tf.AggregationMethod.EXPERIMENTAL_TREE
可能会工作得更好。
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)