Subclassing torch.Tensor

Duncan_Blythe: <p>Hi all,</p>
<p>I would like to subclass torch.Tensor so that my object has some additional data.</p>
<p>It looks something like this:</p>
<pre><code class="lang-auto">class MyObject(torch.Tensor):
    def __init__(self, x, extra_data):
        self.extra_data = extra_data
        super().__init__(x)

obj_ = MyObject([1, 2, 3], extra_data)
</code></pre>
<p>However I get the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "main.py", line 29, in &lt;module&gt;
    sentence = MyObject([1, 2, 3], extra_data)
  File "main.py", line 22, in __init__
    super().__init__(sequence)
TypeError: object.__init__() takes no parameters
</code></pre>

aloukian: <p>I was doing something similar. Because <code>torch.Tensor</code> is a wrapper over a C type, you have to override the <code>__new__</code> method instead. Using your example,</p>
<pre><code class="lang-auto">class MyObject(torch.Tensor):
    @staticmethod
    def __new__(cls, x, extra_data, *args, **kwargs):
        return super().__new__(cls, x, *args, **kwargs)
    
    def __init__(self, x, extra_data):
        #super().__init__() # optional
        self.extra_data = extra_data
        
obj_ = MyObject([1, 2, 3], 'extra_data_123')
obj_.extra_data
</code></pre>
<p>returns <code>'extra_data_123'</code>.</p>
<p>The super call to init is marked optional because I think torch.Tensor uses the default <code>object.__init__</code> call, which is a no-op. However, if you wanted to subclass a child of torch.Tensor that you create in this way, uncomment it.</p>
<p>By the way, you may want to look into the <code>torch.tensor</code> function and put that into the <code>__new__</code> method as well. As is, you cannot pass things like <code>dtype</code> to <code>torch.Tensor</code> or this subclass.</p>

hx2A: <p>I am trying to subclass the torch.Tensor class but am having some difficulties. The above code works great until you call the <code>to</code> or <code>clone</code> methods. Both will return an instance of Tensor, not MyObject.</p>
<p>To improve the above code the best I can accomplish is this:</p>
<pre><code class="lang-auto">class MyObject(torch.Tensor): 
    @staticmethod 
    def __new__(cls, x, extra_data, *args, **kwargs): 
        return super().__new__(cls, x, *args, **kwargs) 
      
    def __init__(self, x, extra_data): 
        self.extra_data = extra_data 

    def clone(self, *args, **kwargs): 
        return MyObject(super().clone(*args, **kwargs), self.extra_data)

    def to(self, *args, **kwargs):
        new_obj = MyObject([], self.extra_data)
        new_obj.data = super().to(*args, **kwargs)
        
        return new_obj
</code></pre>
<p>This works except if <code>requires_grad=True</code>. In that case the <code>to</code> method will detach the object from the graph.</p>
<pre><code class="lang-auto">obj1 = MyObject([1, 2, 3], 'extra_data_123')
obj1.requires_grad_(True)
obj2 = obj1.to('cuda')

obj2.requires_grad
</code></pre>
<p>Returns <code>False</code> but this:</p>
<pre><code class="lang-auto">t1 = torch.Tensor([1, 2, 3])
t1.requires_grad_(True)
t2 = t1.to('cuda')

t2.requires_grad
</code></pre>
<p>Returns <code>True</code>.</p>
<p>How can I improve the <code>to</code> method to make it work like the Tensor class?</p>
<blockquote>
<p>By the way, you may want to look into the  <code>torch.tensor</code>  function and put that into the  <code>__new__</code> method as well. As is, you cannot pass things like  <code>dtype</code>  to  <code>torch.Tensor</code>  or this subclass.</p>
</blockquote>
<p>It would be great if I could find where the <code>torch.tensor</code> function is defined. Can someone tell me?</p>

anantguptadbl: <p>We can override that particular feature specifically. I know it is not exactly subclassing but it can be helpful</p>
<p>Example</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
class MyObject(torch.Tensor): 
    @staticmethod 
    def __new__(cls, x, extra_data, *args, **kwargs): 
        return super().__new__(cls, x, *args, **kwargs) 
      
    def __init__(self, x, extra_data): 
        self.extra_data = extra_data

    def clone(self, *args, **kwargs): 
        return MyObject(super().clone(*args, **kwargs), self.extra_data)

    def to(self, *args, **kwargs):
        new_obj = MyObject([], self.extra_data)
        tempTensor=super().to(*args, **kwargs)
        new_obj.data=tempTensor.data
        new_obj.requires_grad=tempTensor.requires_grad
        return(new_obj)
</code></pre>
<pre><code class="lang-auto">obj1 = MyObject([1, 2, 3], 'extra_data_123')
obj1.requires_grad_(True)
print(obj1.requires_grad)
obj2 = obj1.to('cuda')
print(obj2.requires_grad)
t1 = torch.Tensor([1, 2, 3])
t1.requires_grad_(True)
t2 = t1.to('cuda')
print(t2.requires_grad)
</code></pre>
<pre><code class="lang-auto">True
True
True
</code></pre>

ken012git: <pre><code class="lang-auto">obj1 = MyObject([1], 'extra_data_1')
obj2 = MyObject([2], 'extra_data_2')
print(type(obj1), type(obj2)) # &lt;- &lt;class '__main__.MyObject'&gt; &lt;class '__main__.MyObject'&gt;
sum = obj1 + obj2 
print(sum) # &lt;- tensor([3.])
print(type(sum)) # &lt;- &lt;class '__main__.MyObject'&gt;
obj3 = sum.clone() # &lt;- AttributeError: 'MyObject' object has no attribute 'extra_data'
</code></pre>
<p>I try to fix the issue by adding a class variable:</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
class MyObject(torch.Tensor):
    _extra_data = ''
    
    @staticmethod 
    def __new__(cls, x, extra_data, *args, **kwargs): 
        return super().__new__(cls, x, *args, **kwargs) 
      
    def __init__(self, x, extra_data): 
        self._extra_data = extra_data

    def clone(self, *args, **kwargs): 
        return MyObject(super().clone(*args, **kwargs), self._extra_data)

    def to(self, *args, **kwargs):
        new_obj = MyObject([], self._extra_data)
        tempTensor=super().to(*args, **kwargs)
        new_obj.data=tempTensor.data
        new_obj.requires_grad=tempTensor.requires_grad
        return(new_obj)

    @property
    def extra_data(self):
        return self._extra_data
        
    @extra_data.setter
    def extra_data(self, d):
        self._extra_data = d

obj1 = MyObject([1], 'extra_data_1')
obj2 = MyObject([2], 'extra_data_2')
print(type(obj1), type(obj2)) # &lt;- &lt;class '__main__.MyObject'&gt; &lt;class '__main__.MyObject'&gt;
sum = obj1 + obj2 
print(sum) # &lt;- tensor([3.])
print(type(sum)) # &lt;- &lt;class '__main__.MyObject'&gt;
obj3 = sum.clone() # &lt;- Fine
print(obj3.extra_data) # print nothing
obj3.extra_data = 'obj3_extra_data'
print(obj3.extra_data) # print obj3_extra_data

</code></pre>

Daniel2: <p>There is a bug in the above code that I just spent the last day finding, so I thought I’d share it here. The <code>to</code> function will not propagate gradients properly. This modified <code>to</code> function does do the correct thing</p>
<pre data-code-wrap="python"><code class="lang-python">    def to(self, *args, **kwargs):
        new_obj = super().to(*args, **kwargs)
        if new_obj is self:
            return self
        return MyObject(new_obj, self. extra_data)
</code></pre>