Overwrite forward in layers to work with Data Parallel

username: jsalbert1
cooked: <p>Hello,<br>
I want to create a function to overwrite the forward and backward pass of a nn.Module. E.g. I load a ResNet or any other network, and I automatically change the forward and backward pass of all the layers by a custom Autograd function.</p>
<p>The code I have now is working in CPU and one GPU but not working when I extend to DataParallel.<br>
It is distributing data and the model in different GPUs (Exactly like here: <a href="https://github.com/pytorch/pytorch/issues/8637" class="inline-onebox" rel="noopener nofollow ugc">Issue for DataParallel · Issue #8637 · pytorch/pytorch · GitHub</a>)  and (<a href="https://discuss.pytorch.org/t/dataparallel-on-modules-with-dynamically-overwritten-forwards/29915" class="inline-onebox">DataParallel on modules with dynamically overwritten forwards</a>)</p>
<p>Also I don’t know if that is the best way of overwriting the forward/backward pass.</p>
<p>I define a function to be applied to every model layer:</p>
<pre><code class="lang-auto">def override_backward(layer):
    if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.ConvTranspose2d):
        def forward_conv(x):
            if layer.bias is None:
                return Conv2dFA.apply(x,
                                      layer.weight,
                                      layer.weight_fa,
                                      None,
                                      None,
                                      layer.stride,
                                      layer.padding,
                                      layer.dilation,
                                      layer.groups)
            else:
                return Conv2dFA.apply(x,
                                      layer.weight,
                                      layer.weight_fa,
                                      layer.bias,
                                      layer.bias_fa,
                                      layer.stride,
                                      layer.padding,
                                      layer.dilation,
                                      layer.groups)
        layer.forward = forward_conv
</code></pre>
<p>The function Conv2dFA is:</p>
<pre><code class="lang-auto">class Conv2dFA(autograd.Function):
    @staticmethod
    def forward(context, input, kernels, kernels_fa, bias, bias_fa, stride, padding, dilation, groups):
        context.stride, context.padding, context.dilation, context.groups = stride, padding, dilation, groups
        context.save_for_backward(input, kernels, kernels_fa, bias, bias_fa)
        output = torch.nn.functional.conv2d(input,
                                            kernels,
                                            bias=bias,
                                            stride=stride,
                                            padding=padding,
                                            dilation=dilation,
                                            groups=groups)
        return output

    @staticmethod
    def backward(context, grad_output):
        input, kernels, kernels_fa, bias, bias_fa = context.saved_tensors
        grad_input = grad_kernels = grad_kernels_fa = grad_bias = grad_bias_fa = None

        if context.needs_input_grad[0]:
           grad_input = torch.nn.grad.conv2d_input(input_size=input.shape,
                                                   weight=kernels_fa,
                                                   grad_output=grad_output,
                                                   stride=context.stride,
                                                   padding=context.padding,
                                                   dilation=context.dilation,
                                                   groups=context.groups)

        if context.needs_input_grad[1]:
            grad_kernels = torch.nn.grad.conv2d_weight(input=input,
                                                       weight_size=kernels_fa.shape,
                                                       grad_output=grad_output,
                                                       stride=context.stride,
                                                       padding=context.padding,
                                                       dilation=context.dilation,
                                                       groups=context.groups)

        if bias is not None and context.needs_input_grad[3]:
            grad_bias = grad_output.sum(0).sum(2).sum(1)

        # add the input in the stride gradient which is useless
        # return grad_input, grad_kernels, grad_kernels_fa, grad_bias, grad_bias_fa, stride, padding, dilation, groups
        return grad_input, grad_kernels, grad_kernels_fa, grad_bias, grad_bias_fa, None, None, None, None
</code></pre>
<p>And I apply it like this:</p>
<pre><code class="lang-auto">model_fa = resnet.resnet18()
model_fa.apply(override_backward)
</code></pre>
<p>Is there a way of doing the dynamic overwrite forward that will work with Data Parallel? I don’t want to create custom classes as I want this method to be applicable to every neural net.</p>
<p>Thanks in advance.</p>

username: Yanli_Zhao
cooked: <p>this post showed how to overwrite forward in layers <a href="https://discuss.pytorch.org/t/how-can-i-replace-the-forward-method-of-a-predefined-torchvision-model-with-my-customized-forward-function/54224/6" class="inline-onebox">How can I replace the forward method of a predefined torchvision model with my customized forward function? - #6 by Philipp_Friebertshau</a></p>
<p>hope it helps</p>