Collating named tensors

jteuwen: <p>The named tensor functionality is very promising for the models I use: I work on inverse problems where tensors can have many channels which need to be permuted quite often to fit the operators (such as the torch fft). Having the dimensions named would make life much easier if you e.g. want to sum over a certain named dimension.</p>
<p>However, the API is unstable, and not yet complete. I am looking at what the proper / most elegant way to implement this. I took a rather straightforward approach: in my <code>Dataset</code> class I have transforms which name the dimensions, and when some functions do not support the names, I save them, drop them, apply the operation and add the names again. My datasets typically output dictionaries with keys such as <code>target</code>, <code>image</code> and so on, which are in my case therefore named.</p>
<p>All good. However, when you want to collate these tensors and add a key <code>batch</code> to the first dimension. Unfortunately <code>stack</code> does not yet support named tensors. So I have tried the following, and rewrite the <a href="https://github.com/pytorch/pytorch/blob/00aa23446b9d2b3dac5ed8b343c4536f7d9dd8df/torch/utils/data/_utils/collate.py#L42" rel="nofollow noopener"><code>collate_fn</code></a>:</p>
<pre><code class="lang-python">if isinstance(elem, torch.Tensor):
    out = None
    # TODO: Named tensor: once stack supports named tensors, drop the rename.
    names = elem.names
    elem = elem.rename(None)
    if torch.utils.data.get_worker_info() is not None:
        # If we're in a background process, concatenate directly into a
        # shared memory tensor to avoid an extra copy
        numel = sum([x.rename(None).numel() for x in batch])
        storage = elem.storage()._new_shared(numel)
        out = elem.new(storage)
    out_batch = torch.stack([_.rename(None) for _ in batch], 0, out=out)
    if any(names):
        new_names = tuple(['batch'] + list(names))
        out_batch = out_batch.refine_names(*new_names)
    return out_batch
</code></pre>
<p>Unfortunately that does not work yet.<br>
<code>RuntimeError: pin_memory is not yet supported with named tensors. Please drop names via `tensor = tensor.rename(None)`, call the op with an unnamed tensor, and set names on the result of the operation.</code></p>
<p>The solution here would be to rewrite the DataLoader too, and add this functionality there. If I do this, I am rewriting so much of the core functionality of pytorch (using 1.5) I am wondering: <strong>Is this the intended usage of named tensors, or are they meant to be used in a different place? Is my approach is fine, is there a more convenient way to do all this without having to rewrite everything (e.g. add the names as a key in my dictionary)?</strong></p>

ptrblck: <p>Hi Jonas,</p>
<p>sorry to hear you are trouble with named tensors.<br>
To the best of my knowledge your approach should be correct and these methods are just not implemented yet.</p>
<p>Would the workaround work of returning a <code>dict</code> from your <code>Dataset</code> and add the names afterwards inside the <code>DataLoader</code> loop?</p>

jteuwen: <p>Hi Patrick,</p>
<p>The collate function as I propose above does work if you disable memory pinning which is for my application not a bottleneck (assuming that this will be implemented in the future). You do run into different issues if you do this and have multiple workers:</p>
<blockquote>
<p>RuntimeError: NYI: Named tensors don’t support serialization. Please drop names via <code>tensor = tensor.rename(None)</code> before serialization.</p>
</blockquote>
<p>This is trickier to handle as it is likely somewhere deep in the <code>DataLoader</code>, I will attempt the method of adding names through the dict, will report when and if that works!</p>

jteuwen: <p>This is a solution (hacky) that does work</p>
<pre><code class="lang-auto">class DropNames:
    def __init__(self):
        pass

    def __call__(self, sample):
        new_sample = {}

        for k, v in sample.items():
            if isinstance(v, torch.Tensor) and any(v.names):
                new_sample[k + '_names'] = ';'.join(v.names)  # collate_fn will do funky things without this.
                v = v.rename(None)
            new_sample[k] = v

        return new_sample


class AddNames:
    def __init__(self, add_batch_dimension=True):
        self.add_batch_dimension = add_batch_dimension

    def __call__(self, sample):
        names = [_[:-6] for _ in sample.keys() if _[-5:] == 'names']
        new_sample = {k: v for k, v in sample.items() if k[-5:] != 'names'}

        for name in names:
            if self.add_batch_dimension:
                names = ['batch'] + sample[name + '_names'][0].split(';')

            else:
                names = sample[name + '_names'].split(';')

            new_sample[name] = sample[name].rename(*names)

        return new_sample
</code></pre>
<p>The first function can be added to the end of all the transforms (which are typically composed). This will drop the names, create string out of it separated by <code>;</code> (assuming your variables have no <code>;</code>).<br>
The reason I do this, is because the collate function can do weird things with a list, and if you use the second function as:</p>
<pre><code class="lang-auto">add_names = AddNames(add_batch_dimension=True)
    for iter_idx, data in enumerate(data_loader):
    data = add_names(data)
</code></pre>
<p>you get the names in the tensors back. You might need to drop them again later, as many operations do not  yet support this, but in my case, I make a lot of use of the FFT which in the current implementation requires you to permute the axis often. Having Named Tensors really helps in reducing bugs.</p>
<p>So the above works with pytorch 1.5, so perhaps it is of use to anyone. Improvements are definitely welcome.</p>