How to use dataclass with PyTorch

vainaijr: <p><span class="mention">@dataclass</span><br>
class Net(nn.Module):<br>
…</p>
<p>something like this?</p>

albanD: <p>I don’t use dataclass myself. Do you encounter any problem using it?</p>

vainaijr: <p>dir gives error</p>
<pre><code class="lang-auto">from dataclasses import dataclass

@dataclass
class Net(nn.Module):
  pass  

x = Net()

dir(x)
</code></pre>
<p>gives error</p>
<p><code>AttributeError: 'Net' object has no attribute '_parameters'</code></p>
<p>if I remove <code>@dataclass</code>, then it lists attributes.</p>

ptrblck: <p>It should work, if you initialize the parent class:</p>
<pre><code class="lang-python">@dataclass
class Net(nn.Module):
  def __init__(self):
      super(Net, self).__init__()

x = Net()
dir(x) 
</code></pre>

yassersouri: <p>Although <a class="mention" href="/u/ptrblck">@ptrblck</a>’s answer works, it kind of defeats the purpose of using a dataclass which is partly not writing the <code>__init__</code> function yourself.</p>
<p>So here are some requirements to make this work:</p>
<ol>
<li>The pytorch module class (which is the dataclass itself) needs a <code>__hash__</code> function. The <code>__hash__</code> function is required in the <code>named_modules</code> function of <code>nn.Module</code>.</li>
<li>We need to call <code>super().__init__()</code> at some point.</li>
<li>The dataclass should not be frozen as the <code>__init__</code> function of the <code>nn.Module</code> will try to set attributes. So you cannot use <code>@dataclass(frozen=True)</code> to have a <code>__hash__</code> function for your dataclass.</li>
</ol>
<p>The only solution I found, that is slightly better than <a class="mention" href="/u/ptrblck">@ptrblck</a>’s answer, which I think will work is this:</p>
<pre><code class="lang-python">@dataclass(unsafe_hash=True)
class Net(nn.Module):
    input_feats: int = 10
    output_feats: int = 20
    def __post_init__(self):
        super().__init__()
        self.layer = nn.Linear(self.input_feats, self.output_feats)
</code></pre>
<p>Notice the usage of <code>__post_init__</code> and the ugly hack of setting <code>unsafe_hash=True</code>.</p>

thomas: <p>I wanted to use your solution for my project, however I got an issue when I have a nn.Module as argument:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from dataclasses import dataclass


class Evaluators(nn.Module):
    def __init__(self):
        super(Evaluators, self).__init__()
        self.linear = nn.Linear(1, 1)

@dataclass(unsafe_hash=True)
class Net(nn.Module):
    evaluator: Evaluators
    def __post_init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)

evaluators = Evaluators()
net = Net(evaluators )
</code></pre>
<p>returns:</p>
<pre><code class="lang-auto">test_dataclass.py:18 (test_dataclass)
def test_dataclass():
        evaluators = Evaluators()
&gt;       net = Net(evaluators)

test_dataclass.py:21: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
&lt;string&gt;:3: in __init__
    ???
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = &lt;[ModuleAttributeError("'Net' object has no attribute 'evaluator'") raised in repr()] Net object at 0x17d02449370&gt;
name = 'evaluator'
value = Evaluators(
  (linear): Linear(in_features=1, out_features=1, bias=True)
)

    def __setattr__(self, name: str, value: Union[Tensor, 'Module']) -&gt; None:
        def remove_from(*dicts_or_sets):
            for d in dicts_or_sets:
                if name in d:
                    if isinstance(d, dict):
                        del d[name]
                    else:
                        d.discard(name)
    
        params = self.__dict__.get('_parameters')
        if isinstance(value, Parameter):
            if params is None:
                raise AttributeError(
                    "cannot assign parameters before Module.__init__() call")
            remove_from(self.__dict__, self._buffers, self._modules, self._non_persistent_buffers_set)
            self.register_parameter(name, value)
        elif params is not None and name in params:
            if value is not None:
                raise TypeError("cannot assign '{}' as parameter '{}' "
                                "(torch.nn.Parameter or None expected)"
                                .format(torch.typename(value), name))
            self.register_parameter(name, value)
        else:
            modules = self.__dict__.get('_modules')
            if isinstance(value, Module):
                if modules is None:
&gt;                   raise AttributeError(
                        "cannot assign module before Module.__init__() call")
E                   AttributeError: cannot assign module before Module.__init__() call

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py:807: AttributeError
</code></pre>
<p>Is there a way to solve it, or it just means that I won’t be able to use dataclass ?</p>

yassersouri: <p>I had not considered member variables of type “nn.Module”.<br>
So the solution I proposed will not work in your setting.<br>
I think a solution is possible but requires a bit more hacking into <code>nn.Module</code> or <code>dataclass</code>.</p>

vadimkantorov: <p><a class="mention" href="/u/yassersouri">@yassersouri</a> Please raise your usecase in the issue <a href="https://github.com/pytorch/pytorch/issues/61686" class="inline-onebox" rel="noopener nofollow ugc">[discussion] Remove the need of mandatory super() module call · Issue #61686 · pytorch/pytorch · GitHub</a> <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"> With more comments, maybe the core team would consider it more worthy.</p>

rsokl: <p>I think this can be remedied by the fact that <code>__new__</code> effectively behaves like <code>__pre_init__</code>:</p>
<pre><code class="lang-python">import torch as tr
import torch.nn as nn
from dataclasses import dataclass

@dataclass
class DataclassModule(nn.Module):
    def __new__(cls, *args, **k):
        inst = super().__new__(cls)
        nn.Module.__init__(inst)
        return inst

@dataclass(unsafe_hash=True)
class Net(DataclassModule):
    other_layer: nn.Module
    input_feats: int = 10
    output_feats: int = 20

    def __post_init__(self):
        self.layer = nn.Linear(self.input_feats, self.output_feats)

    def forward(self, x):
        return self.layer(self.other_layer(x))

net = Net(other_layer=nn.Linear(10, 10))
assert net(tr.tensor([1.]*10)).shape == (20,)
assert len(list(net.parameters())) == 4

@dataclass(unsafe_hash=True)
class A(DataclassModule):
    x: int
    def __post_init__(self):
        self.layer1 = nn.Linear(self.x, self.x)

@dataclass(unsafe_hash=True)
class B(A):
    y: int
    def __post_init__(self):
        super().__post_init__()
        self.layer2 = nn.Linear(self.y, self.y)

assert len(list(A(1).parameters())) == 2
assert len(list(B(1, 2).parameters())) == 4
</code></pre>

shinstra: <p>Have had some issues using this solution with nested DataclassModule derived classes. What I observed was some submodules were not transferring weights to the gpu when calling model.cuda(). Digging into it a bit deeper, the parameters of some submodules were not registered as parameters in the  parent modules - which occurred when there were many instances of the same dataclass module in the model.</p>
<p>Using <code>@dataclass(eq=False)</code> instead of  <code>@dataclass(unsafe_hash=True)</code> seems to resolve this. Here is a link to a related discussion:</p>
<p><a href="https://stackoverflow.com/questions/57291307/pytorch-module-with-attrs-cannot-get-parameter-list" rel="noopener nofollow ugc">https://stackoverflow.com/questions/57291307/pytorch-module-with-attrs-cannot-get-parameter-list</a></p>