Concatenate tensors without memory copying

username: lcy
cooked: <p>Hi,</p>
<p>I’m wondering if there is any alternative concatenation method that concatenate two tensor without memory copying?</p>
<p>Currently, I use <code>t = torch.cat([t1, t2], dim=0)</code> in my data pre-processing. However, I got the out-of-memory error because there are many big tensors need to be concatenated.</p>
<p>I have searched around and read some threads like <a href="https://discuss.pytorch.org/t/appending-to-a-tensor/2665">tensor appending</a>, <a href="https://discuss.pytorch.org/t/torch-cat-blows-up-memory-required/8415">Torch.cat() blows up memory required</a>. But still cannot find a desirable solutions to solve the memory consuming problem.</p>

username: bhushans23
cooked: <p>One way is to allocate tensor with target dimension and then assign using slice operator</p>
<p>Following is an example,</p>
<pre><code class="lang-auto">t1_shape = t1.shape[0]
t2_shape = t2.shape[0]
second_dim_shape = t1.shape[1]
t = torch.zeros((t1_shape + t2_shape, second_dim))
t[:t1_shape,:] = t1
t[t1_shape:,:] = t2
</code></pre>
<p>I know that this is a very simple case and will get complicate for nd tensor.<br>
I am not aware of efficient internal api for the same</p>

username: albanD
cooked: <p>Hi,</p>
<p>Since the two original tensors <code>t1</code> and <code>t2</code> are at different places in memory, it’s not possible to make a single Tensor out of them without creating a new tensor that can contain both of them.<br>
The cat implementation does pretty much what the code sample above from <a class="mention" href="/u/bhushans23">@bhushans23</a>  does: Create a Tensor that can contain everything then copy each part into it.</p>

username: lcy
cooked: <p>Ok, so it is inevitable to allocate new memory when concatenate two tensors in pytorch right now.</p>
<p>Is it possible to implement a new concatenation operation like <a href="https://github.com/tensorflow/tensorflow/issues/12948" rel="nofollow noopener">this post</a> in pytorch? It records the reference instead of memory copying.</p>
<p>This new feature should be pretty useful I think.</p>

username: albanD
cooked: <p>Hi,</p>
<p>The post you link does not implement a new concatenation op. It just have a helper function that caries the pair of tensors around and do the ops on each. You can already do that by just using the list with your two Tensors and then adapt the next function to do on it. Did I missed something in that issue?</p>
<p>Another solution is to pre-allocate the full tensor and compute t1 and t2 directly into it doing inplace operations. That way you don’t need the cat operation at all.</p>

username: lcy
cooked: <p>Hi, AlbanD</p>
<p>In that post, the concatenation op doesn’t allocate new memory. It maintains a pointer table which points to the shared memory storage.</p>
<p>The simple solution you suggest below won’t work generally (e.g. the required input type is a tensor rather than a list or I want to concatenate two tensors along with different dimensions).</p>
<blockquote>
<p>You can already do that by just using the list with your two Tensors</p>
</blockquote>
<p>To clarify, let me use a simple example to explain what I want.</p>
<p>Suppose now we concatenate two tensor through below code</p>
<pre><code class="lang-auto">t1 = torch.randn(512, 256, 100, 100)
t2 = torch.randn(512, 256, 100, 100)
t = torch.cat(t1, t2, dim=1) 
</code></pre>
<p>The total memory consuming here will be 512x256x100x100x4 number of float32. Besides, simply list t = [t1, t2] is incorrect.</p>
<p>Is it possible to implement a memory efficient concatenation like</p>
<pre><code class="lang-auto">t1 = torch.randn(512, 256, 100, 100)
t2 = torch.randn(512, 256, 100, 100)
t_efficient = torch.cat(t1, t2, dim=1, allocation="shared")
</code></pre>
<p>The variable t_efficient just records the memory reference of t1 and t2 rather than allocating new memory, and the total memory consuming should be 512x256x100x100x2.</p>

username: albanD
cooked: <p>Hi,</p>
<p>No we don’t have this feature.</p>
<p>What I meant is what this feature does is just keep a list of both Tensors and adapting the few ops you will need to do on them.<br>
A quick example below:</p>
<pre><code class="lang-auto">import torch

class MySharedTensor(object):
    def __init__(self, tensors, dim=None):
        assert dim is not None
        self.dim = dim
        assert (isinstance(tensors, list))
        assert (torch.is_tensor(t) for t in tensors)
        self.tensors = tensors
        self.dim_sizes = [t.size(self.dim) for t in self.tensors]

    # If you want to recover a single Tensor from it
    def to_full_tensor(self):
        return torch.cat(self.tensors, dim=self.dim)

    # Out of place addition
    def __add__(self, other):
        assert torch.is_tensor(other)
        out = other.clone()
        curr_idx = 0
        for i, t in enumerate(self.tensors):
            other_slice = other.narrow(self.dim, curr_idx, self.dim_sizes[i])
            out.narrow(self.dim, curr_idx, self.dim_sizes[i]).copy_(t).add_(other_slice)
            curr_idx += self.dim_sizes[i]
        return out

    # Inplace add
    def __iadd__(self, other):
        assert torch.is_tensor(other)
        curr_idx = 0
        for i, t in enumerate(self.tensors):
            other_slice = other.narrow(self.dim, curr_idx, self.dim_sizes[i])
            t.add_(other_slice)
            curr_idx += self.dim_sizes[i]
        return self

    # Matrix Multiplication (only 2d matrices for simplicity)
    def mm(self, other):
        assert other.ndimension() == 2
        assert all(t.ndimension() == 2 for t in self.tensors)

        if self.dim == 0:
            out_tensors = []
            for t in self.tensors:
                out_tensors.append(t.mm(other))
            return MySharedTensor(out_tensors, dim=0)
        elif self.dim == 1:
            out = 0
            curr_idx = 0
            for i, t in enumerate(self.tensors):
                other_slice = other.narrow(0, curr_idx, self.dim_sizes[i])
                out += t.mm(other_slice)
                curr_idx += self.dim_sizes[i]
            return out
        else:
            raise RuntimeError("Invalid dimension")


a = torch.rand(2, 4)
b = torch.rand(2, 4)
print("a and b")
print(a)
print(b)

c = MySharedTensor([a, b], dim=1)
print("c size:")
print(c.to_full_tensor().size())

d = torch.rand(2, 8)
print("d")
print(d)

e = c + d
print("e = c + d")
print(e)

f = c.to_full_tensor() + d
print("f = c.to_full_tensor() + d")
print(f)

c += d
print("c += d")
print(c.to_full_tensor())
print("a and b are changed:")
print(a, b)


g = torch.rand(8, 3)
print("g")
print(g)

h = c.mm(g)
print("h = c.mm(g)")
print(h)

k = c.to_full_tensor().mm(g)
print("k = c.to_full_tensor().mm(g)")
print(k)
</code></pre>

username: lcy
cooked: <p>The MySharedTensor is a temporary solution because it only works in a few of ops and cannot seamlessly adapt to all ops that the native Tensor supports. This is the reason I said it won’t work generally.</p>
<p>I hope the pytorch team can add this new feature in the future.</p>
<p>By the way, I want MySharedTensor can work with conv2d without allocating new memory. What should I do?</p>

username: albanD
cooked: <p>The thing is that supporting such thing is a very large change as each op needs to be reimplemented for it. In particular, if you don’t want to ever allocate the full Tensor, it can be tricky to use libraries like cudnn that do not support these features without a speed drop.<br>
Do you know how tensorflow handle such cases? What is the performance drop of using such structure?</p>
<p>For the case of conv2d, that will depend on which dimension you concatenate over.<br>
If it’s 0, then you can do a for-loop passing each Tensor one by one and then concatenating the outputs along dimension 0.<br>
If it’s 1, then you will need 2 smaller convs (if you have 2 input Tensors) that take a subset of the input channels and the same number of output channels and sum the output of each conv.<br>
If it’s 2 or 3 then it gets quite tricky to do as the interface between the two is shared and you cannot do two independant conv.</p>

username: lcy
cooked: <p>Sorry, I don’t know how tensorflow deal with this case. But it seems a technique report <a href="https://arxiv.org/abs/1707.06990" rel="nofollow noopener">Memory-Efficient Implementation of DenseNets</a> has already implemented this feature in pytorch.</p>
<p>I will concatenate only in dimension 0 and 1. Will try this trick. Thanks for your help.</p>

username: jnhwkim
cooked: <p>For the newcomers, the implementation of DenseNets uses <a href="https://pytorch.org/docs/1.2.0/checkpoint.html?highlight=checkpoint" rel="noopener nofollow ugc">checkpoint</a>, where intermediate features (concatenation) are recovered by re-forwarding for each segment while back-warding.</p>
<blockquote>
<p>Checkpointing is implemented by rerunning a forward-pass segment for each checkpointed segment during backward. This can cause persistent states like the RNG state to be advanced than they would without checkpointing</p>
</blockquote>

username: albanD
cooked: <p>By the way this is coming to pytorch <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"><br>
See this issue: <a href="https://github.com/pytorch/pytorch/issues/22169" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/22169</a></p>

username: 1ce9b3de5641ab6f51c6
cooked: <p>But what if I can guarantee that concatenated tensors are views of the same tensor. Imagine I have a tensor</p>
<pre><code class="lang-auto">t = torch.randn(400)
</code></pre>
<p>and I need to calculate the mean of every 100 elements. In order to avoid any loops, I can easily do something like this:</p>
<pre><code class="lang-auto">means = t.view(4, -1).mean(dim=1)
</code></pre>
<p>but let’s say I need to calculate <code>mean()</code> of intersecting segments, in particular:</p>
<pre><code class="lang-auto">means = torch.cat((t[0:200], t[100:300], t[200:400]), dim=0).mean(dim=1)
</code></pre>
<p>Can I do it efficiently, without tensor copy happening in <code>torch.cat()</code>?</p>

username: albanD
cooked: <p>You can in some cases yes, but it won’t be the most easy-to-read code:</p>
<pre><code class="lang-python">import torch

a = torch.rand(40)

print(a)

size_slices = 20
size_overlap = 10
nb_slices = 3
b = a.as_strided((nb_slices, size_slices), (size_slices-size_overlap, 1), 0)
print(b)

print(b.sum(0))
</code></pre>

username: Pingo
cooked: <p>Hi guys,<br>
I have been tracking this topic for a long time, but still couldn’t find a good solution.<br>
So, I tried to implement a customized pytorch cuda extension to totally avoid memory copy on cuda device. However, my solution is very specific in use case.</p>
<p>In my case, I combined Concat + Conv2d into a single CatConv2d kernel, which can significantly reduce the latency in some cases (small batch, small in/out channels).</p>
<p>for example, the operations of concat of a list of Tensor followed by a conv2d before, now can be done in a single operator with single cuda kernel:</p>
<pre><code class="lang-auto">    x_list = [x0, x1, x2, x3,...]
    x = torch.cat(x_list, 1)
    x = conv.forward(x)

    # equivalent to 
    x = CatConv2d(x_list)
</code></pre>
<p>Here is the repo:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/PingoLH/CatConv2d">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/PingoLH/CatConv2d" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/928e82cf80102ae9b54b79a6157e7784ee2b482a06411491157c2186946a8ae1/PingoLH/CatConv2d" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/PingoLH/CatConv2d" target="_blank" rel="noopener nofollow ugc">GitHub - PingoLH/CatConv2d: Concat+Conv2d fused all-in-one CUDA kernel...</a></h3>

  <p>Concat+Conv2d fused all-in-one CUDA kernel extension for Pytorch - GitHub - PingoLH/CatConv2d: Concat+Conv2d fused all-in-one CUDA kernel extension for Pytorch</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>The backward path is still missing (forward only) in this repo, and there are many limitations, but I’m just trying to give a sense that concat+some_op combination is very promising in terms of performance and worth to work on.</p>

username: sparshgarg23
cooked: <aside class="quote no-group" data-username="lcy" data-post="6" data-topic="34609">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/ecccb3/48.png" class="avatar"> lcy:</div>
<blockquote>
<pre><code class="lang-auto">t1 = torch.randn(512, 256, 100, 100)
t2 = torch.randn(512, 256, 100, 100)
t = torch.cat(t1, t2, dim=1) 
</code></pre>
</blockquote>
</aside>
<p>I ran your code on colab,and got the following error<br>
cat() received an invalid combination of arguments - got (Tensor, Tensor, dim=int), but expected one of:</p>
<ul>
<li>(tuple of Tensors tensors, int dim, *, Tensor out)</li>
<li>(tuple of Tensors tensors, name dim, *, Tensor out)</li>
</ul>

username: ptrblck
cooked: <p>Pass the tensors as a <code>tuple</code> or <code>list</code> as described in the error message:</p>
<pre><code class="lang-python">t = torch.cat((t1, t2), dim=1) 
t = torch.cat([t1, t2], dim=1) 
</code></pre>

username: dbsx
cooked: <p>I am wondering if this concatenate without copying could be done in C++?<br>
<a href="https://discuss.pytorch.org/t/how-to-concatenate-tensors-without-copy/151146">How to concatenate tensors without copy? - C++ - PyTorch Forums</a></p>