How to work with `functional_call`? parametrization of model

bpfrd: <p>Hi,<br>
I want to parametrize my model to accept <code>params</code> in the forward. If it is given, the model uses these params otherwise the model uses the registered parameters (default). The params have the same structure as <code>model.named_parameters()</code>.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision.models as models
from torch.func import functional_call

class ParameterizedModel(nn.Module):
    def __init__(self, embedding_dim=64):
        super(ParameterizedModel, self).__init__()
        self.base_model = models.resnet18(pretrained=True)
        self.base_model.fc = nn.Linear(self.base_model.fc.in_features, embedding_dim)

    def forward(self, x, params=None):
        if params is not None:
            return functional_call(self.base_model, params, x)
        else:
            return self.base_model(x)

model = ParameterizedModel(embedding_dim=64)
model
&gt;&gt;&gt;
ParameterizedModel(
  (base_model): ResNet(
    (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu): ReLU(inplace=True)
    (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
    (layer1): Sequential(
      (0): BasicBlock(
        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
      (1): BasicBlock(
        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (layer2): Sequential(
      (0): BasicBlock(
        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (downsample): Sequential(
          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): BasicBlock(
        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (layer3): Sequential(
      (0): BasicBlock(
        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (downsample): Sequential(
          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): BasicBlock(
        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (layer4): Sequential(
      (0): BasicBlock(
        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (downsample): Sequential(
          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): BasicBlock(
        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
    (fc): Linear(in_features=512, out_features=64, bias=True)
  )
)
</code></pre>
<p>Looking at the instructions of <code>functional_call</code> in <a href="https://pytorch.org/docs/stable/generated/torch.func.functional_call.html" class="inline-onebox" rel="noopener nofollow ugc">torch.func.functional_call — PyTorch 2.3 documentation</a> I did the following:</p>
<pre><code class="lang-auto">from collections import OrderedDict

params = OrderedDict()
for name, m in model.named_parameters():
    l, r = name.rsplit('.', 1)
    new_name = f'{l}.parametrizations.{r}.original'
    params[new_name] = m

x = torch.rand(1, 3, 256, 256)
model(x, params=params)
</code></pre>
<p>But I get the below error:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/torch/nn/utils/_named_member_accessor.py in get_submodule(self, name)
    135         try:
--&gt; 136             return self.memo[name]
    137         except KeyError:

KeyError: 'base_model.conv1.parametrizations.weight'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
18 frames
KeyError: 'base_model.conv1.parametrizations'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
KeyError: 'base_model.conv1'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
KeyError: 'base_model'

During handling of the above exception, another exception occurred:

AttributeError                            Traceback (most recent call last)
AttributeError: 'ResNet' object has no attribute 'base_model'

The above exception was the direct cause of the following exception:

AttributeError                            Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/torch/nn/utils/_named_member_accessor.py in get_submodule(self, name)
    144                 submodule = getattr(module, attr)
    145             except AttributeError as ex:
--&gt; 146                 raise AttributeError(
    147                     f"{module._get_name()} has no attribute `{attr}`"
    148                 ) from ex

AttributeError: ResNet has no attribute `base_model`
</code></pre>
<p>How can I achieve this goal for any model without hard-coding this for each model?</p>
<p>best</p>

bpfrd: <p>I also tried this workaround:</p>
<pre><code class="lang-auto">import torch.nn.functional as F

def flatten_module(module, prefix=''):
    _ = []
    for name, child in module.named_children():
        child_prefix = f'{prefix}.{name}' if prefix else name
        if isinstance(child, nn.Module):
            _.extend(flatten_module(child, child_prefix))
        else:
            _.append((child_prefix, child))
    if not list(module.named_children()):
        _.append((prefix, module))
    return _

def apply_layer(layer, x, params=None, name=''):
    if isinstance(layer, nn.Conv2d):
        weight = params.get(name + '.weight', layer.weight)
        bias = params.get(name + '.bias', layer.bias)
        x = F.conv2d(x, weight, bias, layer.stride, layer.padding, layer.dilation, layer.groups)
    elif isinstance(layer, nn.Linear):
        weight = params.get(name + '.weight', layer.weight)
        bias = params.get(name + '.bias', layer.bias)
        x = F.linear(x, weight, bias)
    elif isinstance(layer, nn.BatchNorm2d):
        weight = params.get(name + '.weight', layer.weight)
        bias = params.get(name + '.bias', layer.bias)
        running_mean = layer.running_mean
        running_var = layer.running_var
        x = F.batch_norm(x, running_mean, running_var, weight, bias, layer.training, layer.momentum, layer.eps)
    elif isinstance(layer, nn.ReLU):
        x = F.relu(x, inplace=layer.inplace)
    elif isinstance(layer, nn.MaxPool2d):
        x = F.max_pool2d(x, layer.kernel_size, layer.stride, layer.padding, layer.dilation, layer.ceil_mode)
    elif isinstance(layer, nn.AdaptiveAvgPool2d):
        x = F.adaptive_avg_pool2d(x, layer.output_size)
    elif isinstance(layer, nn.Dropout):
        x = F.dropout(x, layer.p, layer.training, layer.inplace)
    else:
        x = layer(x)
    return x
</code></pre>
<p>here’s the output (structure is similar to <code>model.named_parameters()</code>):</p>
<pre><code class="lang-auto">flattened = flatten_module(model)
flattened
&gt;&gt;
[('base_model.conv1',
  Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)),
 ('base_model.bn1',
  BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.relu', ReLU(inplace=True)),
 ('base_model.maxpool',
  MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)),
 ('base_model.layer1.0.conv1',
  Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer1.0.bn1',
  BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer1.0.relu', ReLU(inplace=True)),
 ('base_model.layer1.0.conv2',
  Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer1.0.bn2',
  BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer1.1.conv1',
  Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer1.1.bn1',
  BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer1.1.relu', ReLU(inplace=True)),
 ('base_model.layer1.1.conv2',
  Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer1.1.bn2',
  BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer2.0.conv1',
  Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)),
 ('base_model.layer2.0.bn1',
  BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer2.0.relu', ReLU(inplace=True)),
 ('base_model.layer2.0.conv2',
  Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer2.0.bn2',
  BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer2.0.downsample.0',
  Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)),
 ('base_model.layer2.0.downsample.1',
  BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer2.1.conv1',
  Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer2.1.bn1',
  BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer2.1.relu', ReLU(inplace=True)),
 ('base_model.layer2.1.conv2',
  Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer2.1.bn2',
  BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer3.0.conv1',
  Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)),
 ('base_model.layer3.0.bn1',
  BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer3.0.relu', ReLU(inplace=True)),
 ('base_model.layer3.0.conv2',
  Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer3.0.bn2',
  BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer3.0.downsample.0',
  Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)),
 ('base_model.layer3.0.downsample.1',
  BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer3.1.conv1',
  Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer3.1.bn1',
  BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer3.1.relu', ReLU(inplace=True)),
 ('base_model.layer3.1.conv2',
  Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer3.1.bn2',
  BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer4.0.conv1',
  Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)),
 ('base_model.layer4.0.bn1',
  BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer4.0.relu', ReLU(inplace=True)),
 ('base_model.layer4.0.conv2',
  Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer4.0.bn2',
  BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer4.0.downsample.0',
  Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)),
 ('base_model.layer4.0.downsample.1',
  BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer4.1.conv1',
  Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer4.1.bn1',
  BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.layer4.1.relu', ReLU(inplace=True)),
 ('base_model.layer4.1.conv2',
  Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
 ('base_model.layer4.1.bn2',
  BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
 ('base_model.avgpool', AdaptiveAvgPool2d(output_size=(1, 1))),
 ('base_model.fc', Linear(in_features=512, out_features=64, bias=True))]
</code></pre>
<p>and</p>
<pre><code class="lang-auto">x = torch.rand(1, 3, 256, 256)
params = dict(model.named_parameters())
for (name,layer) in flattened:
    x = apply_layer(layer, x, params, name)
</code></pre>
<p>I get this error (probably because I don’t follow the order of operations in the <code>base_model.forward</code> function:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-11-d14d92ed2d0f&gt; in &lt;cell line: 32&gt;()
     31 params = dict(model.named_parameters())
     32 for (name,layer) in flattened:
---&gt; 33     x = apply_layer(layer, x, params, name)

&lt;ipython-input-11-d14d92ed2d0f&gt; in apply_layer(layer, x, params, name)
      5         weight = params.get(name + '.weight', layer.weight)
      6         bias = params.get(name + '.bias', layer.bias)
----&gt; 7         x = F.conv2d(x, weight, bias, layer.stride, layer.padding, layer.dilation, layer.groups)
      8     elif isinstance(layer, nn.Linear):
      9         weight = params.get(name + '.weight', layer.weight)

RuntimeError: Given groups=1, weight of size [128, 64, 1, 1], expected input[1, 128, 32, 32] to have 64 channels, but got 128 channels instead
</code></pre>
<p>does anyone have any idea?</p>