TensorFlow:如何衡量每个张量占用多少 GPU 内存? [英] TensorFlow: How to measure how much GPU memory each tensor takes?

问题描述

我目前正在 TensorFlow 中实现 YOLO 并且我有点惊讶于占用了多少内存.在我的 GPU 上，我可以使用他们的 Darknet 框架训练 YOLO，批量大小为 64.在 TensorFlow 上我只能用批量大小为 6，我已经用完了 8 的内存.对于测试阶段，我可以在不耗尽内存的情况下以 64 的批处理大小运行.

I'm currently implementing YOLO in TensorFlow and I'm a little surprised on how much memory that is taking. On my GPU I can train YOLO using their Darknet framework with batch size 64. On TensorFlow I can only do it with batch size 6, with 8 I already run out of memory. For the test phase I can run with batch size 64 without running out of memory.

1. 我想知道如何计算每个张量消耗了多少内存?默认情况下所有张量都保存在 GPU 中吗?我可以简单地将总内存消耗计算为形状 * 32 位吗?

1. I am wondering how I can calculate how much memory is being consumed by each tensor? Are all tensors by default saved in the GPU? Can I simply calculate the total memory consumption as the shape * 32 bits?

我注意到由于我使用的是动量，所以我所有的张量也有一个 /Momentum 张量.这也会占用大量内存吗?

I noticed that since I'm using momentum, all my tensors also have a /Momentum tensor. Could that also be using a lot of memory?

我正在使用一种方法 distorted_inputs 扩充我的数据集，该方法与 CIFAR-10 教程.莫非这部分占用了大量内存?我相信暗网会在 CPU 中进行修改.

I am augmenting my dataset with a method distorted_inputs, very similar to the one defined in the CIFAR-10 tutorial. Could it be that this part is occupying a huge chunk of memory? I believe Darknet does the modifications in the CPU.

推荐答案

现在 1258 已经关闭，可以在导入 TensorFlow 之前通过设置环境变量来启用 Python 中的内存日志记录:

Now that 1258 has been closed, you can enable memory logging in Python by setting an environment variable before importing TensorFlow:

import os
os.environ['TF_CPP_MIN_VLOG_LEVEL']='3'
import tensorflow as tf

因此将有大量的日志记录.您需要对结果进行 grep 以找到合适的行.例如:

There will be a lot of logging as a result of this. You'll want to grep the results to find the appropriate lines. For example:

grep MemoryLogTensorAllocation train.log

这篇关于TensorFlow:如何衡量每个张量占用多少 GPU 内存?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！