如何解决联合Tensorflow内存-重载单个GPU
我正在尝试使用张量流进行联合学习。
我已经基于this tutorial for MNIST创建了一个模型。现在,我对其进行更改:我为其提供尺寸为112x112的图像,每个图像具有3个通道(即输入层的尺寸为112x112x3)。当我尝试使用多个模型(大约50个,每个大约100张图像)时,出现“内存不足”异常。查看--log-file
的输出是有道理的:由于某种原因,我的程序仅从单个GPU占用内存。我该如何避免呢?
一些评论:
- 我使用
nvidia-smi
,所以我很确定显示的内存使用情况是实际的内存使用情况。我确实可以看到它随着时间的推移而增长。 - TensorFlow可以识别其他GPU:在我运行程序之前,它们已占用0MB内存。日志中也有第
allow_growth
行。
解决方法
您在TFF中遇到了一个有趣而棘手的问题。
TLDR:TFF以非正统的方式使用TF;这碰到了TF中的错误,其中设备放置无法正确地通过tf.data.Dataset.reduce
传播,而TFF的操作仅放置在单个GPU上。
TFF贡献者当前正在为TFF的默认配置添加一种解决方法,该方法实际上涉及或多或少的“拉动” TensorFlow的AutoGraph来生成tf.while
循环,而不是减少数据集。如果您是从源代码构建的,则可以将dataset.reduce
交换为特定的for循环:
num_examples_sum = 0
for batch in iter(dataset):
num_examples_sum = reduce_fn(num_examples_sum,batch)
这将导致所有GPU被利用。
如果它们仍在放大,请尝试向上调整clients_per_thread
参数(与先前num_client_executors
参数相反); TFF默认会尝试并行运行所有客户端。
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。