Subclassing torch.Tensor

1

Hi all,

I would like to subclass torch.Tensor so that my object has some additional data.

It looks something like this:

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)

However I get the following error:

Traceback (most recent call last):
  File "main.py", line 29, in <module>
    sentence = MyObject([1, 2, 3], extra_data)
  File "main.py", line 22, in __init__
    super().__init__(sequence)
TypeError: object.__init__() takes no parameters
2

I was doing something similar. Because torch.Tensor is a wrapper over a C type, you have to override the __new__ method instead. Using your example,

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

returns 'extra_data_123'.

The super call to init is marked optional because I think torch.Tensor uses the default object.__init__ 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.

By the way, you may want to look into the torch.tensor function and put that into the __new__ method as well. As is, you cannot pass things like dtype to torch.Tensor or this subclass.

3

I am trying to subclass the torch.Tensor class but am having some difficulties. The above code works great until you call the to or clone methods. Both will return an instance of Tensor, not MyObject.

To improve the above code the best I can accomplish is this:

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

This works except if requires_grad=True. In that case the to method will detach the object from the graph.

obj1 = MyObject([1, 2, 3], 'extra_data_123')
obj1.requires_grad_(True)
obj2 = obj1.to('cuda')

obj2.requires_grad

Returns False but this:

t1 = torch.Tensor([1, 2, 3])
t1.requires_grad_(True)
t2 = t1.to('cuda')

t2.requires_grad

Returns True.

How can I improve the to method to make it work like the Tensor class?

By the way, you may want to look into the torch.tensor function and put that into the __new__ method as well. As is, you cannot pass things like dtype to torch.Tensor or this subclass.

It would be great if I could find where the torch.tensor function is defined. Can someone tell me?

4

We can override that particular feature specifically. I know it is not exactly subclassing but it can be helpful

Example

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)

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)

True
True
True
5 
obj1 = MyObject([1], 'extra_data_1')
obj2 = MyObject([2], 'extra_data_2')
print(type(obj1), type(obj2)) # <- <class '__main__.MyObject'> <class '__main__.MyObject'>
sum = obj1 + obj2 
print(sum) # <- tensor([3.])
print(type(sum)) # <- <class '__main__.MyObject'>
obj3 = sum.clone() # <- AttributeError: 'MyObject' object has no attribute 'extra_data'

I try to fix the issue by adding a class variable:

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)) # <- <class '__main__.MyObject'> <class '__main__.MyObject'>
sum = obj1 + obj2 
print(sum) # <- tensor([3.])
print(type(sum)) # <- <class '__main__.MyObject'>
obj3 = sum.clone() # <- Fine
print(obj3.extra_data) # print nothing
obj3.extra_data = 'obj3_extra_data'
print(obj3.extra_data) # print obj3_extra_data
6

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 to function will not propagate gradients properly. This modified to function does do the correct thing

    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)