.to(dtype) method with custom module and buffer of fixed precision

Azerus: <p>Hello,</p>
<p>I made a custom module that <strong>need</strong> high precision (float64) for few operations in the forward pass. Those operations involve few float64 tensors registered as buffer in the module as shown in the following snipped:</p>
<pre><code class="lang-auto">class custom_module(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('foo_tensor', get_foo_tensor())
        self.register_buffer('float64_tensor',  get_float64_tensor())

    def forward(self, x):
        input_dtype = x.dtype
        x = some_regular_op(x, self.foo_tensor)

        # Avoid any autocast context messing around here
        with torch.cuda.amp.autocast(enabled=False):
            x = x.to(dtype = torch.float64)
            x = some_high_precision_op(x, self.float64_tensor)

        x = x.to(input_dtype)
        return x
</code></pre>
<p>As I registered my tensors, I can easily change the device of my module with <code>.to(device)</code>. But now, I would like also to be able to use <code>.to(dtype)</code>, but without affecting the precision of <code>float64_tensor</code>, is that possible in some way ?<br>
Furthermore, my <code>custom_module</code> may be used as sub-module inside a model, so I want to be able to perform <code>model.to(dtype)</code> and still preserve the precision of this specific <code>float64_tensor</code>.</p>
<p>Any help would be great <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>If you think that’s a good idea too, maybe I can ask for a freeze_dtype option for the <code>register_buffer</code> method as feature request.</p>

Azerus: <p>So, the simpler workaround I found is to avoid registering my float64_tensor, and move it to the current device if needed during the forward pass. Plus I added hooks to the <code>state_dict</code> in order to correctly set or get this <code>float64_tensor</code> when saving or loading my module.</p>
<p>Snipped of code with workaround:</p>
<pre><code class="lang-auto">class custom_module(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('foo_tensor', get_foo_tensor())
        float64_tensor = get_float64_tensor()

        # Register state_dict hooks to get/set float64_tensor when saving/loading the module
        def set_f64_tensor_from_state_dict_hook(state_dict, *args, **kwargs):
                self.float64_tensor = state_dict['float64_tensor']
        self._register_load_state_dict_pre_hook(set_f64_tensor_from_state_dict_hook)
        def update_state_dict_hook(self, state_dict, *args, **kwargs):
                state_dict['float64_tensor'] = self.float64_tensor
                return state_dict
        self._register_state_dict_hook(update_state_dict_hook)

    def forward(self, x):
        input_dtype = x.dtype
        x = some_regular_op(x, self.foo_tensor)

        # Change device of f64 tensor if needed
        self.float64_tensor = self.float64_tensor.to(x.device)

        # Avoid any autocast context messing around here
        with torch.cuda.amp.autocast(enabled=False):
            x = x.to(dtype = torch.float64)
            x = some_high_precision_op(x, self.float64_tensor)

        x = x.to(input_dtype)
        return x
</code></pre>
<p>Which is a good enough solution to me. But please, don’t hesitate to comment if you see any problem with this approach or if you have any smarter move.</p>

tom: <aside class="quote no-group" data-username="Azerus" data-post="2" data-topic="126990">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/azerus/48/33745_2.png" class="avatar"> Azerus:</div>
<blockquote>
<p>So, the simpler workaround I found is to avoid registering my float64_tensor, and move it to the current device if needed during the forward pass. Plus I added hooks to the <code>state_dict</code> in order to correctly set or get this <code>float64_tensor</code> when saving or loading my module.</p>
</blockquote>
</aside>
<p>I think you found a pretty good (and thorough) solution here, thank you for sharing!<br>
You might double check if it works with distributed if that is something you need, and it seems inside the documented API. In particular, the obvious alternative – override <code>Module.to</code> –  would also work but bear more risk from messing with PyTorch internals.</p>
<p>Best regards</p>
<p>Thomas</p>

Azerus: <p>Thanks for the feedback.</p>
<p>Good point about the distributed. I’m not using it, but I like to think of most use case (unfortunatly, I have only one gpu so i don’t know if I can actually perform the test).</p>
<p>I also though about overwriting <code>Module.to</code> but it would not work if my <code>custom_module</code> is registered as sub-module of model when running <code>model.to(dtype)</code>, as only the <code>to()</code> method from <code>model</code> would be called and the not the one from <code>custom_module</code>. To do such approach I would have to overwrite the <code>.to</code> method from the <code>float64_tensor</code> itself, if I understood the <a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module.to" rel="noopener nofollow ugc">source</a> correctly (see the line <code>return self._apply(convert)</code> and how the function <code>convert</code> is defined).</p>