How can multiple tensors conveniently be stored in one object?

pmeier: <p>For a project of mine I need to bundle some <code>torch.Tensor</code>s into one object, since they have a strong connection. Think something like this:</p>
<pre><code class="lang-python">import torch

a = torch.zeros(5)
b = torch.ones(4)
c = torch.rand(3)

class MyObject:
    def __init__(self, a, b, c):
        self.a = a
        self.b = b
        self.c = c

my_object = MyObject(a, b, c)
</code></pre>
<p><code>MyObject</code> is only used for convenient storage (think <code>namedtuple</code>). I have two requirements for <code>MyObject</code>:</p>
<ol>
<li>Regarding the utility methods, <code>MyObject</code> should act like a single <code>torch.Tensor</code>. For example <code>my_object.cuda()</code> should move all stored <code>torch.Tensor</code>s to the GPU.</li>
<li>If embedded into a <code>torch.nn.Module</code> it should be treated like another <code>torch.Tensor</code>. Thus, if I call <code>state_dict()</code> it should also extract the <code>torch.Tensor</code>s stored in <code>MyObject</code>.</li>
</ol>
<p>The only way I see to achieve this is is by making <code>MyObject</code> an <code>nn.Module</code>:</p>
<pre><code class="lang-python">from torch import nn

class TensorStorage(nn.Module):
    def __init__(self, **attrs):
        super().__init__()
        for name, attr in attrs.items():
            if isinstance(attr, torch.Tensor):
                self.register_buffer(name, attr)
            else:
                setattr(self, name, attr)

    def forward(self):
        msg = (f"{self.__class__.__name__} objects are only used "
                "for storage and cannot be called.")
        raise RuntimeError(msg)

my_object = TensorStorage(a=a, b=b, c=c)
</code></pre>
<p>As far as I’ve tested this works for my purpose. However, being an <code>nn.Module</code> implies that it can be called like every other module. Although this can be caught with an error in the <code>forward()</code> method, it might be misleading for the user. My Questions are:</p>
<ol>
<li>Can someone think of a better approach to address my requirements?</li>
<li>If not, can I use <code>TensorStorage</code> as implemented or will I get in trouble somewhere down the road?</li>
</ol>

cpuhrsch: <p>Hello Philip,</p>
<p>If you’re interested in ongoing development we’re currently working on <a href="https://github.com/pytorch/nestedtensor" rel="nofollow noopener">https://github.com/pytorch/nestedtensor</a> as a first-class representation of collections of variably sized Tensors. I’m happy to talk further about this if you want and would love your feedback.</p>
<p>Thanks,<br>
Christian</p>

pmeier: <p>Hi Christian,</p>
<p>at first sight this looks exactly like what I imagined. Thanks for pointing me towards this. Unfortunately I need to finish my project first and I decided to go with the quirky work around at first. Afterwards I’m happy to try to achieve the same with <code>nestedtensor</code>. How would I get in touch for feedback or suggestions? Should I use this forum, a GitHub issue or something else entirely?</p>

cpuhrsch: <p>Hello Philip,</p>
<p>Ideally you’d create issues on the nestedtensor Github repo for wider requests, bugs etc. But I’ll also see your responses to this thread in case something is not aligned with those buckets of topics.</p>
<p>Thanks for giving it a shot! Just to set expectations: at this point the repository is still under heavy development so they will be the occasional bug or change in semantics and general sparsity in documentation. I expect to improve on this very quickly however.</p>
<p>Thanks,<br>
Christian</p>