How can I convert specific continuous layers?

username: yukkyo
cooked: <p>I am currently trying to implement the following paper. And I defined these modules.<br>
<a href="https://arxiv.org/abs/1911.09737" rel="nofollow noopener">[1911.09737] Filter Response Normalization Layer: Eliminating Batch Dependence in the Training of Deep Neural Networks</a></p>
<p>Next, I want to convert <code>BatchNorm2d + ReLU</code> to <code>FRN + TLU</code> in my network.<br>
(ex. <code>ReLU</code> not after <code>BatchNorm2d</code> is excluded)<br>
First, I used bellow <code>bnrelu_to_frn()</code>. However, this method does not consider forward order.</p>
<p>How can I convert specific continuous layers ?</p>
<pre><code class="lang-auto">class TLU(nn.Module):
    def __init__(self, num_features):
        super(TLU, self).__init__()
        self.num_features = num_features
        self.tau = nn.parameter.Parameter(
            torch.Tensor(1, num_features, 1, 1), requires_grad=True)

    def reset_parameters(self):
        nn.init.zeros_(self.tau)

    def forward(self, x):
        return torch.max(x, self.tau)


class FRN(nn.Module):
    def __init__(self, num_features, eps=1e-6, is_eps_leanable=False):
        super(FRN, self).__init__()

        self.num_features = num_features
        self.init_eps = eps
        self.is_eps_leanable = is_eps_leanable

        self.weight = nn.parameter.Parameter(
            torch.Tensor(1, num_features, 1, 1), requires_grad=True)
        self.bias = nn.parameter.Parameter(
            torch.Tensor(1, num_features, 1, 1), requires_grad=True)
        if is_eps_leanable:
            self.eps = nn.parameter.Parameter(torch.Tensor(1), requires_grad=True)
        else:
            self.register_buffer('eps', torch.Tensor([eps]))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)
        if self.is_eps_leanable:
            nn.init.constant_(self.eps, self.init_eps)

    def forward(self, x):
        nu2 = x.pow(2).mean(dim=[2, 3], keepdim=True)
        x = x * torch.rsqrt(nu2 + self.eps.abs())
        x = self.weight * x + self.bias
        return x

def bnrelu_to_frn(module):
    mod = module
    before_name = None
    before_child = None
    is_before_bn = False

    for name, child in module.named_children():
        if is_before_bn and isinstance(child, (ReLU, LeakyReLU)):
            if isinstance(before_child, BatchNorm2d):
                mod.add_module(
                    before_name, FRN(num_features=before_child.num_features))
            else:
                raise NotImplementedError()
            mod.add_module(name, TLU(num_features=before_child.num_features))
        else:
            mod.add_module(name, bnrelu_to_frn(child))

        before_name = name
        before_child = child
        is_before_bn = isinstance(child, BatchNorm2d)
    return mod
</code></pre>

username: mobassir94
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> any solution for this?</p>

username: ptrblck
cooked: <p>As explained by the author of this question, the execution order is not maintained by using <code>.named_children()</code> and modules will be returned based on the initialization. I don’t know if there is a clean way to perform this type of model surgery automatically without e.g. overriding some modules manually.</p>