[QUESTION] feature fusion and re-input

[bpfliegel]

I am generally fine (and quite happy) about feature extraction of models as per: https://rwightman.github.io/pytorch-image-models/feature_extraction/

However, I can't figure out how to solve something generally given any backbone used:

Think of a multimodal semantic segmentation task, assuming we have RGB and D as inputs. Both RGB and D is sent through their respective encoder branches (one backbone for each), we fuse features at stride 2,4,8,etc. and apply some decoder. This approach is very easy to do.

Okay. now imagine that we have a third encoder branch, which does not have any input, but gets the fusion product from the RGB and D branches at stride 2, and having that as its stride 2 input, invokes all modules to arrive at stride 4. At stride 4 it fuses its features with the stride 4 features of RGB and D branches and arrives at stride 8, etc. etc.

As I am aware TIMM generally assumes that we have an input at the start of the network and we parametrize TIMM to give us the outputs at specific strides. I am afraid there is no such feature to say, that 'hey, this is the input for stride 2, give me output at stride 4'. Or is there a way, but I failed to find that? I would be really useful to support fusion architectures.

I would welcome any ideas on how to do that, if it's possible.

Thanks a lot,
Balint

[rwightman]

@bpfliegel there is no pre-set feature for that...

You can use the feature extraction functionality in a general purpose way if you determine the other points of interest for extraction you can use the module names to get features anywhere.

Injecting inputs back into the network is something else enterirely and would be a significant complexity to support across models. Probably the only way to try and do this without modifying significant parts of a model, would be to use FX to rewrite the model based on the points where you want to inject/fuse and extract

[rwightman]

Moving to discussions as this is not a bug, and outside scope of a feature request at this time...

[bpfliegel]

Thanks for your kind answer @rwightman, really appreciate it. I will check on torch FX to see if I can do that! Thanks so much again, Balint

[bpfliegel]

Torch.FX has limitations, and some networks have their own unique madness - but this works most of time if someone needs it:

import torch
import numpy as np
import torch.nn as nn
from torch.fx import symbolic_trace
import timm
from typing import Any, Callable, Dict, Optional, Tuple
 
# https://github.com/pytorch/examples/blob/main/fx/module_tracer.py
class ModulePathTracer(torch.fx.Tracer):
    current_module_qualified_name : str = ''
    node_to_originating_module : Dict[torch.fx.Node, str] = {}
 
    def call_module(self, m: torch.nn.Module, forward: Callable[..., Any],
                    args : Tuple[Any, ...], kwargs : Dict[str, Any]) -> Any:
        old_qualname = self.current_module_qualified_name
        try:
            self.current_module_qualified_name = self.path_of_module(m)
            return super().call_module(m, forward, args, kwargs)
        finally:
            self.current_module_qualified_name = old_qualname
 
    def create_proxy(self, kind: str, target: torch.fx.node.Target, args: Tuple[Any, ...],
                     kwargs: Dict[str, Any], name: Optional[str] = None, type_expr: Optional[Any] = None):
        proxy = super().create_proxy(kind, target, args, kwargs, name, type_expr)
        self.node_to_originating_module[proxy.node] = self.current_module_qualified_name
        return proxy
 
class ModularTimmNetwork(nn.Module):
    def __init__(self, backbone_name, in_channels):
        super().__init__()
       
        # Create model
        self.model = timm.create_model(backbone_name, features_only=True, exportable=True, pretrained=True, in_chans=in_channels)
 
        # Get output features, channels, strides
        fi = self.model.feature_info
        out_indices = fi.out_indices
        self.features = []
        for feature_index, feature in enumerate(fi.info):
            if feature_index in out_indices:
                self.features.append(feature)
 
        # Symbolic trace - tracing also originating module
        tracer = ModulePathTracer()
        fx_model = tracer.trace(self.model)
 
        # Record exit index for all features
        exits = []
        for feature in self.features:
            module_name = feature['module']
            last_seen = 0
            for node_index, node in enumerate(fx_model.nodes):
                module_qualname = tracer.node_to_originating_module.get(node)
                if not (module_qualname is None):
                    if module_qualname.startswith(module_name):
                        last_seen = node_index
            exits.append(last_seen)
 
        # Create layers from subgraphs
        self.layers = nn.ModuleList()
        for feature_index, feature in enumerate(self.features):
            module_name = feature['module']
            start_index = (0 if (feature_index==0) else exits[feature_index-1]) + 1
            end_index = exits[feature_index]
           
            new_graph = torch.fx.Graph()
            value_remap = {}
            nodes_to_remove = []
 
            new_input_node = new_graph.placeholder('x')
            old_input_name = None
 
            for node_index, node in enumerate(fx_model.nodes):
                if (node_index<=end_index):
                    new_node = new_graph.node_copy(node, lambda n : value_remap[n])
                    value_remap[node] = new_node
                    if (node_index == start_index):
                        assert(len(new_node.all_input_nodes)==1)
                        old_input_name = new_node.all_input_nodes[0].name
                    if (node_index < start_index):
                        nodes_to_remove.append(new_node)
                    if (node_index>=start_index):
                        for input_node in new_node.all_input_nodes:
                            if (input_node.name == old_input_name):
                                new_node.replace_input_with(input_node, new_input_node)
 
            # Add last node as output
            new_graph.output(new_node)
 
            # Erase non-needed nodes in reverse order
            for n in reversed(nodes_to_remove):
                new_graph.erase_node(n)
 
            # Lint
            new_graph.lint()
 
            # Create module based on subgraph
            self.layers.append(torch.fx.GraphModule(self.model, new_graph))
 
    def forward(self, x):
        return self.model(x)
 
    def forward_by_layers(self, x):
        outs = []
        for layer_index, layer in enumerate(self.layers):
            out = layer(x)
            outs.append(out)
            x = out
        return outs
 
    def forward_layer(self, x, layer_index):
        return self.layers[layer_index](x)
 
def main():
    net = ModularTimmNetwork('efficientnet_b0', 3)
    net.eval()
    dummy_input = torch.FloatTensor((np.random.rand(1,3,256,256) * 2.0) - 1.0)
    result1 = net.forward(dummy_input)
    result2 = net.forward_by_layers(dummy_input)
    result3_0 = net.forward_layer(dummy_input, 0)
    result3_1 = net.forward_layer(result3_0, 1)
    result3_2 = net.forward_layer(result3_1, 2)
    result3_3 = net.forward_layer(result3_2, 3)
    check1 = result1[3].detach().numpy()
    check2 = result3_3.detach().numpy()
    print(np.array_equal(check1, check2))
 
if __name__ == '__main__':
    main()

Quick and dirty solution, sorry for that. Thanks for the idea @rwightman on torch.fx!