Torch.max(tensor1, tensor2) memory problem

Zhang_Fuyang: <p>Hi,</p>
<p>For example, if I have bunch of feature maps, let’s say we have 100 feature map in the shape [128, 256, 256].<br>
then I want to get the pixel-wise max feature map of this 100 feature.<br>
the easiest way is to use torch.cat() to concatenate the total 100 feature map then use torch.max() to get the result. But my memory is not big enough to store the tensor from concatenation.<br>
So, I think maybe I can use for-loop to get the max tensor. code is like this:</p>
<pre><code class="lang-auto">for i in range(100):
    if i == 0:
        out = feature[i]
    else:
        out = torch.max(out, feature[i])
</code></pre>
<p>but I find that memory is still growing. I believe it is because this torch.max(out, feature[i]) also need to create a continuous space to hold the tensor. But is there a way to free this memory after the calculation? Because ideally we don’t need this memory to do auto-grad, we only need to remember each pixels comes from then we can get gradient.</p>

ptrblck: <p>I assume <code>feature</code> is a list holding all tensors?<br>
If that’s the case, you would have enough memory to at least store the data once.<br>
Could you try to create a <code>feature</code> tensor in the first place instead of a list?<br>
If you are working close to the memory limit, the memory peaks might yield an OOM error.</p>
<p>That being said, if you don’t need gradients for the calculation, you could wrap the code in a <code>with torch.no_grad()</code> block.</p>

Zhang_Fuyang: <p>Thanks for replying.</p>
<p>I find that my statement is misleading. Sorry about that, let me make the problem more clear.</p>
<p>Suppose I have 300 vector store in a tensor called <code>feature</code>, which is in the shape [300, 1024]. And we also have an <code>index</code> tensor, which is in [300] size and indicates which vector should send into max() operation to get maximum tensor. so, <code>index</code> is something like [1,0,1,0,0,1,…,0].</p>
<p>The easiest way to do this is just one line code:<br>
<code>out = torch.max(feature[torch.where(index == 1)], 0)</code><br>
But, we have a limit memory, and I guess this part <code>feature[torch.where(index == 1)]</code> will allocate a new memory, which is not acceptable.</p>
<p>So, I think there is another way: first, allocate space in memory for storage maximum result. Then, use a for-loop to compare each feature with result one-by-one.</p>
<pre><code class="lang-auto">out = torch.zeros([1024])
for i in range(300):
    if index[i] == 1:
        out = torch.max(out, feature[i])
</code></pre>
<p>So, ideally, we only need to allocate one [1024] size memory for <code>out</code> tensor, and this is acceptable by our gpu. (Both <code>index</code> and <code>feature</code> are tensors here, but we don’t want to create a big intermediate variable so we use for-loop)</p>
<p>However, I find that during the for-loop, each iteration, memory is still growing. I don’t know why, and how to avoid it.</p>
<p>We need  <code>out</code> tensor contains gradient for back-propagation. So I guess this is the result pytorch doesn’t release the memory after <code>torch.max(out, feature[i])</code>. But, ideally, we just need one same size tensor with <code>out</code> to remember where the max value come from, then we can do backpropagation. Is there a way to do that?</p>