Extracting GraphNode Layers from Pretrained Network and converting them to nn.Module format.

[Parikshit00]

I selected my backbone as:
backbone = caformer.CAFormerS18(input_shape=(3,256, 256), pretrained="imagenet", num_classes = 0)
then i have some layers names i am interested in: layer_names = ['stack4_block3_mlp_Dense_1', 'stack3_block9_mlp_Dense_1', 'stack2_block3_mlp_Dense_1', 'stack1_block3_mlp_Dense_1']
then I select these layers from backbone: selected_layers = [backbone.get_layer(x).output for x in layer_names]
now forexample: I initialize first layer x = layers[0] and then i want to apply some nn.Module upscale_feature = decode(x, scale = 4, filters = x.shape[channel_axis])

i get error: TypeError: unsupported operand type(s) for %: 'GraphNode' and 'int'
my decoder:

class decode(nn.Module):
    def __init__(self, input_channels, filters, scale=2, activation='relu'):
        super(decode, self).__init__()
        self.conv1 = nn.Conv2d(input_channels, filters, kernel_size=(1, 1), padding="same", bias=False)
        nn.init.kaiming_uniform_(self.conv1.weight, mode='fan_in', nonlinearity='relu')
        self.activation1 = nn.ReLU()
        etc.....

    def forward(self, x):
        x = self.conv1(x)
        x = self.activation1(x)
         return x

is there a method that convers 'GraphNode' to the default nn.Module() torch type objects so i can connect these layers to my custom Torch nn.Module classes? or is it that i should rewrite my Torch classes to inherit only methods from keras_cv_attention_models.pytorch_backend for eg: x1 = layers.Conv2D()(input_tensor)instead of nn.Conv2d(input_channels, filters, kernel_size=(1, 1), padding="same", bias=False). Thanks.

[leondgarse]

I'm glad you are interested in trying this. :)

• Shortly in your usage case:

  import torch
  from torch import nn
  from keras_cv_attention_models import caformer
  
  backbone = caformer.CAFormerS18(input_shape=(3, 256, 256), pretrained="imagenet", num_classes=0)
  layer_names = ['stack4_block3_mlp_Dense_1', 'stack3_block9_mlp_Dense_1', 'stack2_block3_mlp_Dense_1', 'stack1_block3_mlp_Dense_1']
  selected_layers = [backbone.get_layer(x) for x in layer_names]  # Don't use .output
  
  class Decode(nn.Module):
      def __init__(self, input_channels, filters=32, scale=2, activation='relu'):
          super().__init__()
          self.dense1 = nn.Linear(input_channels, filters, bias=False)
          nn.init.kaiming_uniform_(self.dense1.weight, mode='fan_in', nonlinearity='relu')
          self.activation1 = nn.ReLU()
  
      def forward(self, x):
          x = self.dense1(x)
          x = self.activation1(x)
          return x
  
  aa = selected_layers[0]
  bb = Decode(aa.output.shape[-1])
  print(bb(aa(torch.ones([1, *aa.input_shape[1:]]))).shape)
  # torch.Size([1, 64, 32])

• Most time the values retured using model.get_layer are layers defined in layers.py, like Dense Conv2D. Here Dense itself is already a torch.nn.Module, and Dense.module is the raw nn.Linear.

  print(aa.__class__.__name__)
  # Dense
  print(isinstance(aa, torch.nn.Module))
  # True
  print(aa.module)
  # Linear(in_features=2048, out_features=512, bias=False)

• The xxx.output is a GraphNode, points to the node itself. Just a mimic behavior like keras.layers.output, recording things like output shape, dtype, etc.

  print(aa.output.shape, aa.output.dtype)
  # [None, 64, 512] torch.float32
  print(aa.output.layer is aa)
  # True

[Parikshit00]

@leondgarse Thanks. Also, this can also be done by registering forward hooks to the desired layers and detaching the activations since the loaded backbone is nn.Module so native torch methods also works for this but i get other issues with shape mismatch with keras-cv-attention-models-1.3.24 i get a different shape for the layers 'stack4_block3_mlp_Dense_1', 'stack3_block9_mlp_Dense_1' than the expected shape. expected: [<KerasTensor: shape=(None, 8, 8, 512) dtype=float32 (created by layer 'stack4_block3_mlp_Dense_1')>, <KerasTensor: shape=(None, 16, 16, 320) dtype=float32 (created by layer 'stack3_block9_mlp_Dense_1')>, <KerasTensor: shape=(None, 32, 32, 128) dtype=float32 (created by layer 'stack2_block3_mlp_Dense_1')>, <KerasTensor: shape=(None, 64, 64, 64) dtype=float32 (created by layer 'stack1_block3_mlp_Dense_1')>] , And I got (with tf type call)[<KerasTensor: shape=(None, 64, 512) dtype=float32 (created by layer 'stack4_block3_mlp_Dense_1')>, <KerasTensor: shape=(None, 256, 320) dtype=float32 (created by layer 'stack3_block9_mlp_Dense_1')>, <KerasTensor: shape=(None, 32, 32, 128) dtype=float32 (created by layer 'stack2_block3_mlp_Dense_1')>, <KerasTensor: shape=(None, 64, 64, 64) dtype=float32 (created by layer 'stack1_block3_mlp_Dense_1')>]...........I solved it with keras-cv-attention-models-1.3.9 ON TF backend.......HOWEVER, on using Torch backend on 1.3.9.... the tensors arent Torch type and neither is the backbone model and the shape mismatch arises..... Maybe torch backend wasnot supported yet with 1.3.9?. In this case how should i proceed that the Torch backend works and I get proper shape with 1.3.24 version?

[leondgarse]

• The output shape changing 4D -> 3D is introduced in update attention blocks using 3D input for better inference performance. Just some of my testing results show that 3D inputs perform better. Thus you may try versions before that like kecam==1.3.17.
• Using the newest 1.3.24, comment off those lines in caformer.py in the above check-in if need 4D outputs. Just some Reshape layers.
• Another issue in my previous reply is that we need to create a model from inputs till the layers that are required, not using the layer itself. Usage instance could be:

  import torch
  from torch import nn
  from keras_cv_attention_models import caformer
  from keras_cv_attention_models.backend import models
  
  backbone = caformer.CAFormerS18(input_shape=(3, 256, 256), pretrained="imagenet", num_classes=0)
  layer_names = ['stack4_block3_mlp_Dense_1', 'stack3_block9_mlp_Dense_1', 'stack2_block3_mlp_Dense_1', 'stack1_block3_mlp_Dense_1']
  
  # Create a new model from inputs till layers required
  new_backbone = models.Model(backbone.inputs, [backbone.get_layer(ii).output for ii in layer_names])
  print([ii.shape for ii in new_backbone(torch.ones([1, 3, 256, 256]))])
  # [torch.Size([1, 8, 8, 512]), torch.Size([1, 16, 16, 320]), torch.Size([1, 32, 32, 128]), torch.Size([1, 64, 64, 64])]
  
  # Using the new backbone in a torch module
  class Decode(nn.Module):
      def __init__(self, filters=32, scale=2, activation='relu'):
          super().__init__()
          self.backbone = new_backbone
          self.denses = nn.Sequential(*[nn.Linear(ii[-1], filters, bias=False) for ii in self.backbone.output_shape])
          # nn.init.kaiming_uniform_(self.dense1.weight, mode='fan_in', nonlinearity='relu')
          self.activation1 = nn.ReLU()
  
      def forward(self, x):
          features = self.backbone(x)
          features = [self.activation1(dd(ff)) for ff, dd in zip(features, self.denses)]
          return features
  
  aa = Decode()
  print([ii.shape for ii in aa(torch.ones([1, 3, 256, 256]))])
  # [torch.Size([1, 8, 8, 32]), torch.Size([1, 16, 16, 32]), torch.Size([1, 32, 32, 32]), torch.Size([1, 64, 64, 32])]

[Parikshit00]

@leondgarse Thanks...it works. Sorry for late reply :)