Body: <h2>I'm trying to implement a model guided by the paper of <a href="https://arxiv.org/abs/2306.09626" rel="nofollow noreferrer">PAtt-lite</a> (weights are on <a href="https://github.com/jlrex/patt-lite" rel="nofollow noreferrer">https://github.com/jlrex/patt-lite</a> but no implementation provided yet). Using FER2013+ and RAF-DB. The main classes I've implemented are:
  <a href="https://i.sstatic.net/fq1Wc.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/fq1Wc.png" alt="enter image description here" /></a></h2>
  <p>Classes for separable convolutional layers: Depthwise and pointwise.</p>
  <pre><code>import math
  import torch
  import torch.nn.functional as F
  from transformers import MobileNetV1Config, MobileNetV1Model, AutoImageProcessor
  
  
  class DepthWiseSeperable(nn.Module):
  
      def __init__(self, in_channels, out_channels, stride, padding, kernel_size):
  
          super(DepthWiseSeperable, self).__init__()
  
          # groups used here
          self.depthwise = nn.Conv2d(in_channels = in_channels , out_channels = in_channels , stride = stride , padding = padding, kernel_size = kernel_size , groups=in_channels , bias = False)
          self.relu = nn.PReLU(in_channels)
          self.bn = nn.BatchNorm2d(in_channels)
  
      def forward(self, x):
  
          x = self.depthwise(x)
          x = self.relu(x)
          x = self.bn(x)
          return x
  
  
  class PointWiseSeperable(nn.Module):
  
      def __init__(self, in_channels, out_channels, stride):
  
          super(PointWiseSeperable, self).__init__()
  
          self.pointwise =  nn.Conv2d(in_channels = in_channels , out_channels = out_channels , stride = 1 , padding = 0, kernel_size = 1, bias = False)
          self.relu = nn.PReLU(out_channels)
          self.bn = nn.BatchNorm2d(out_channels)
  
      def forward(self, x):
  
          x = self.pointwise(x)
          x = self.relu(x)
          x = self.bn(x)
          return x
  
  
  class SepConv2d(nn.Module):
      def __init__(self, in_channels, out_channels, stride, padding, kernel_size):
  
          super(SepConv2d, self).__init__()
  
          self.depthwise = DepthWiseSeperable(in_channels, in_channels, stride, padding, kernel_size)
          self.pointwise = PointWiseSeperable(in_channels, out_channels, stride)
  
  
      def forward(self, x):
          x = self.depthwise(x)
          x = self.pointwise(x)
  
          return x
  </code></pre>
  <hr />
  <p>Since later on we are going to create an instance of MobileNetV1 from huggingface a class is implemented to truncate such net. Paper says:</p>
  <blockquote>
  <p>layers after the depthwise convolution of block 9 are skipped</p>
  </blockquote>
  <pre><code>class Truncated_MobileNet(nn.Module):
      
      def __init__(self, original_model):
          super(Truncated_MobileNet, self).__init__()
          self.conv_stem = original_model.conv_stem
          self.layer = nn.ModuleList(original_model.layer[:17])  #layers after the depthwise convolution of block 9 are skipped
  
      def forward(self, x):
          x = self.conv_stem(x)
          for layer in self.layer:
              x = layer(x)
    
          return x
  
  </code></pre>
  <hr />
  <p>Here patch extraction module. In the paper they say that pading is needed from 14x14 to 16x16, but I already obtain 16x16 from truncated mobilenet. I adjust the kernel sizes and strides to match dimension with the architecture graphic.</p>
  <pre><code>class PatchExtraction(nn.Module):
  
      def __init__(self, in_channels, out_channels):
          &quot;&quot;&quot;
          Patch extraction block.
  
          Args:
              in_channels (int): Number of input channels.
              out_channels (int): Number of output channels.
          &quot;&quot;&quot;
          super(PatchExtraction, self).__init__()
  
          # Padding to achieve 16x16 spatial dimensions
          self.padding = nn.ZeroPad2d((1, 1, 1, 1))
  
          # Depthwise separable convolution output 4x4x256
          self.depthwise1 = SepConv2d(in_channels = 512, out_channels = 256, stride = 4 , padding = 0, kernel_size = 4)
          # Depthwise separable convolution output 2x2x256
          self.depthwise2 = SepConv2d(in_channels = 256, out_channels = 256, stride = 2 , padding = 0, kernel_size = 2)
  
          # Pointwise convolution output 2x2x256
          self.pointwise = PointWiseSeperable(in_channels = 256, out_channels = 256, stride=1)
  
  
  
      def forward(self, x):
  
          x = self.depthwise1(x)
          x = self.depthwise2(x)
          x = self.pointwise(x)
  
          return x
  </code></pre>
  <hr />
  <p>The attention classifier module. <strong>I am using residual as I saw that stabilises the training, but I don't know if this way is the best practice.</strong></p>
  <pre><code>class AttentionClassifier(nn.Module):
      def __init__(self, in_features, hidden_size, num_classes, d_q):
          super(AttentionClassifier, self).__init__()
          
          # First fully connected layer
          self.fc1 = nn.Linear(in_features, hidden_size)
          self.relu = nn.ReLU()
  
          # Linear projections for query, key, and value
          self.fc_query = nn.Linear(hidden_size, d_q)
          self.fc_key = nn.Linear(hidden_size, d_q)
          self.fc_value = nn.Linear(hidden_size, d_q)  
  
      
          # Second fully connected layer
          self.fc2 = nn.Linear(d_q, num_classes)
          #torch.nn.init.xavier_uniform_(self.fc2.weight)
          self.softmax = nn.Softmax(dim=-1)
          # ScaledDotProductAttention instance
          #self.attention = ScaledDotProductAttention(d_q)
  
      def forward(self, x):
          
          # Apply the first fully connected layer
          y = self.fc1(x)
          y = self.relu(y)
          # Linear projections for query, key, and value
          query = self.fc_query(y)
          key = self.fc_key(y)
          value = self.fc_value(y)
  
          attention_output = F.scaled_dot_product_attention(query, key, value, dropout_p= 0.0, scale = None)
          attention_output = attention_output + x
          # Apply the second fully connected layer
          x = self.softmax(self.fc2(attention_output))
          
          return x
  </code></pre>
  <hr />
  <p>Joining patch extraction and attention modules. <strong>I don't know if input_dim&lt;hidden_dim is the best use of attention.</strong></p>
  <pre><code>class Top_model(nn.Module):
      def __init__(self,  num_classes=7):
          
          super(Top_model, self).__init__()
          
          self.patchextraction = PatchExtraction(512, 256)
          
          # Global Average Pooling and classifier
  
          self.classifier = nn.Sequential(
              AttentionClassifier(in_features=256, hidden_size=512, num_classes=7, d_q=256),
          )
          self.gap = nn.AdaptiveAvgPool2d((1,1))
  
          
      def forward(self, x):
          x = self.patchextraction(x)
          x = self.gap(x)
          x = x.view(x.size(0), -1)  # Flattening the tensor
          x = self.classifier(x)
  
          return x    
  </code></pre>
  <hr />
  <p>Finally join all:</p>
  <pre><code>class Patt_lite(nn.Module):
      
      def __init__(self, truncated_model, top_model, freeze_base=True, freeze_top = True):
          
          super(Patt_lite, self).__init__()
  
          self.base = truncated_model
          if freeze_base:
              for i, param in enumerate(self.base.parameters()):
                  param.requires_grad = False
          
          self.top = top_model
          if freeze_top:
              for i, param in enumerate(self.top.parameters()):
                  param.requires_grad = False
          
      def forward(self, x):
          x = self.base(x)
          x = self.top(x)
  
          return x      
  </code></pre>
  <hr />
  <p>For the training I'm using Fastai:</p>
  <pre><code># Data augmentation for training
  train_batch_transforms = [
       # Quality
       Brightness(p=0), Contrast(p=0), Saturation(p=0, draw = 0),
       # Data Alterations
       RandomErasing(p=0.25),
       # Orientation
       Flip(p=0.5), Rotate(p=0.75, draw=30), Warp(p=0.5),
       #Norm
       Normalize.from_stats([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
   ]
  train_item_transforms = [
      Resize(224),
      ToTensor(),
  ]
  
  test_batch_transforms = [
       Normalize.from_stats([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
   ]
  test_item_transforms = [
      #ToTensor, Normalize
      Resize(224),
      ToTensor(),
      
  ]
  
  #Set batch size
  batch_size = 8
  
  
  # Create data loaders using the defined transforms
  dls = ImageDataLoaders.from_folder(
      path='path to train dataset',
      valid_pct=0.2,
      bs=batch_size,
      valid_bs=batch_size,
      item_tfms=train_item_transforms,
      batch_tfms=train_batch_transforms,
      device=&quot;cuda&quot;,
      convert_mode='L'
  )
  
  #Model instance
  configuration = MobileNetV1Config(tf_padding=False)
  original_model = MobileNetV1Model.from_pretrained(&quot;Matthijs/mobilenet_v1_1.0_224&quot;, config=configuration)
  truncated_model = Truncated_MobileNet(original_model)
  
  top_model = Top_model(num_classes = 7)
  patt_model = Patt_lite(truncated_model, top_model, freeze_base = False, freeze_top = False)
  patt_model = patt_model.to(device)
  
  #Optimization and learning
  opt_func = partial(OptimWrapper, opt=optim.RMSprop)
  learn = Learner(dls, patt_model, opt_func=opt_func, metrics=accuracy)
  learn.loss_func = CrossEntropyLossFlat()
  learn.fit(n_epoch=50, lr = 1e-4, cbs = [ShowGraphCallback()])
  </code></pre>
  <p>I tried <code>Matthijs/mobilenet_v1_1.0_224</code> and <code>google/mobilenet_v1_1.0_224</code>, but seems the same.</p>
  <p>On the paper they split the training in two parts: top model training with mobilenet freezed and fine tuning of mobilenet. <strong>Are the weights of mobilenet correctly loaded?</strong></p>
  <p>The first training result in bad performance and the second one do not improve the accuracy or loss.</p>
  <p><a href="https://i.sstatic.net/CzdrC.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/CzdrC.png" alt="enter image description here" /></a>
  <a href="https://i.sstatic.net/hqTz8.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/hqTz8.png" alt="enter image description here" /></a></p>
  <p>Any idea? Thanks</p>
  <hr />
  <p>Edit:
  Option A: Eliminate dot_product_attention (also softmax and residual), so</p>
  <pre><code>class AttentionClassifier(nn.Module):
      def __init__(self, in_features, hidden_size, num_classes, d_q):
          super(AttentionClassifier, self).__init__()
          
          # First fully connected layer
          self.fc1 = nn.Linear(in_features, hidden_size)
          self.relu = nn.ReLU()
  
     
      
          # Second fully connected layer
          self.fc2 = nn.Linear(d_q, num_classes)
  
      def forward(self, x):
          
          # Apply the first fully connected layer
          x = self.fc1(x)
          x = self.relu(x)
         
          # Apply the second fully connected layer
          x = self.fc2(x)
          #x = self.softmax(x)
          
          return x
  </code></pre>
  <p>After 25 epoch:
  <a href="https://i.sstatic.net/cLhlD.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/cLhlD.png" alt="enter image description here" /></a>
  <a href="https://i.sstatic.net/ywW93.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/ywW93.png" alt="enter image description here" /></a></p>
  <hr />
  <p>Option B: No patch, no dot_product_attention, just MobileNet (freezed) and 2 FC layers.
  <a href="https://i.sstatic.net/Qd0ke.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/Qd0ke.png" alt="enter image description here" /></a>
  <a href="https://i.sstatic.net/WMtc1.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/WMtc1.png" alt="enter image description here" /></a></p>
  <hr />
  <p>Option C: same as Option B but with mobilenet unfreezed. Clearly overfitting.
  <a href="https://i.sstatic.net/zbUYo.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/zbUYo.png" alt="enter image description here" /></a>
  <a href="https://i.sstatic.net/zll0R.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/zll0R.png" alt="enter image description here" /></a></p>
  <hr />
  <p>Option D: No patch extraction, attention without residual.</p>
  <pre><code>class AttentionClassifier(nn.Module):
      def __init__(self, in_features, hidden_size, num_classes, d_q):
          super(AttentionClassifier, self).__init__()
          
          # First fully connected layer
          self.fc1 = nn.Linear(in_features, hidden_size)
          self.relu = nn.ReLU()
  
          # Linear projections for query, key, and value
          self.fc_query = nn.Linear(hidden_size, d_q)
          self.fc_key = nn.Linear(hidden_size, d_q)
          self.fc_value = nn.Linear(hidden_size, d_q)  
  
      
          # Second fully connected layer
          self.fc2 = nn.Linear(d_q, num_classes)
  
      def forward(self, x):
          
          # Apply the first fully connected laye
          x = self.fc1(x)
          x = self.relu(x)
          # Linear projections for query, key, and value
          query = self.fc_query(x)
          key = self.fc_key(x)
          value = self.fc_value(x)
  
          attention_output = F.scaled_dot_product_attention(query, key, value, dropout_p= 0.0, scale = None)
          #Apply the second fully connected layer
          x = self.fc2(attention_output)
          #x = self.softmax(x)
          
          return x
  
  
      
  class Top_model(nn.Module):
      
      def __init__(self,  num_classes=7):
          
          super(Top_model, self).__init__()
          
          self.patchextraction = PatchExtraction(512, 256)
          
          # Global Average Pooling and classifier
  
          self.classifier = nn.Sequential(
              AttentionClassifier(in_features=512, hidden_size=256, num_classes=num_classes, d_q=256),
          )
          self.gap = nn.AdaptiveAvgPool2d((1,1))
  
          
      def forward(self, x):
          x = self.gap(x)
          x = x.view(x.size(0), -1)  # Flattening the tensor
          x = self.classifier(x)
  
          return x   
  </code></pre>
  <p><a href="https://i.sstatic.net/11jFG.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/11jFG.png" alt="enter image description here" /></a>
  <a href="https://i.sstatic.net/5zOPL.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/5zOPL.png" alt="enter image description here" /></a></p>
  <hr />
  <p>Option E: Same as option D but with <code> attention_output = attention_output + x</code></p>
  <p><a href="https://i.sstatic.net/scc8F.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/scc8F.png" alt="enter image description here" /></a>
  <a href="https://i.sstatic.net/n0j5F.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/n0j5F.png" alt="enter image description here" /></a></p>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T13:45:14.073Z
Comments.Text: In my opinion it's always important to make sure that the model trains first. 
  
  Since your training loss does not really decrease, it's interesting to know if you have tried changing the hyperparameters? For example increasing the learning rate? 
  
  I also noticed that in your attention layer you are adding x as a residual element and not the y, that looks like an input to attention head. Is it intentional?
Comments.UserDisplayName: vl_knd
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T13:48:25.570Z
Comments.Text: I don't see any problems in the code at first glance, but what can be useful, if you have in fact checked that its not hyperparametetrs problem, is to try simpler base or top model. For top you can just flatten all the features and apply one fully connected layer for example. Like this you can assure yourself that specific parts of your model in fact are working as they are suppose to. You can really do this for all the important "elements" of the model, where the problem can be hidden in your opinion.
Comments.UserDisplayName: vl_knd
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T13:58:59.287Z
Comments.Text: The using of x and y was intentional since i tried to increase the hidden dimension, so the hidden dimension was grater than k,q, v and attention output dimension, thus y = self.fc1(x) does not match with attention_output but x matches. I tried also with input_dim = hidden_dim = output_dim and use y as a residual: same performance.
Comments.UserDisplayName: Robert
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T14:00:03.643Z
Comments.Text: Thanks! I will try it part by part to address the issue.
Comments.UserDisplayName: Robert
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T14:03:22.827Z
Comments.Text: @vl_knd Do i need a flatten between GlobalAvgPool and the classifier aprt in top model forward? It is enough with the AdaptativeAvgPool2d((1,1)) and x = x.view(x.size(0), -1)??
Comments.UserDisplayName: Robert
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T14:25:33.223Z
Comments.Text: I would suggest you replace x.view(x.size(0), -1) by torch.flatten (https://pytorch.org/docs/stable/generated/torch.flatten.html)
Comments.UserDisplayName: vl_knd
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T14:26:15.350Z
Comments.Text: Can you try doing global average polling on feature maps of mobile net and training your attention classifier first, to see if it works? Without the "patch extraction block"?
Comments.UserDisplayName: vl_knd
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-02-13T23:03:04.513Z
Comments.Text: Let us [continue this discussion in chat](https://chat.stackexchange.com/rooms/151468/discussion-between-vl-knd-and-robert).
Comments.UserDisplayName: vl_knd
Tags: <computer-vision><attention><transfer-learning><facial-recognition><state-of-the-art>
Title: SOTA model (PAtt-lite) implementation with Pytorch
UserDisplayName: Robert
ThreadText: # SOTA model (PAtt-lite) implementation with Pytorch
  
  (Asked by: username_0 on 2024-02-13)
  
  ## I'm trying to implement a model guided by the paper of [PAtt-lite](https://arxiv.org/abs/2306.09626) (weights are on <https://github.com/jlrex/patt-lite> but no implementation provided yet). Using FER2013+ and RAF-DB. The main classes I've implemented are: [enter image description here](https://i.sstatic.net/fq1Wc.png)
  
  Classes for separable convolutional layers: Depthwise and pointwise.
  
  ```
  import math
  import torch
  import torch.nn.functional as F
  from transformers import MobileNetV1Config, MobileNetV1Model, AutoImageProcessor
  
  class DepthWiseSeperable(nn.Module):
  
      def __init__(self, in_channels, out_channels, stride, padding, kernel_size):
  
          super(DepthWiseSeperable, self).__init__()
  
          # groups used here
          self.depthwise = nn.Conv2d(in_channels = in_channels , out_channels = in_channels , stride = stride , padding = padding, kernel_size = kernel_size , groups=in_channels , bias = False)
          self.relu = nn.PReLU(in_channels)
          self.bn = nn.BatchNorm2d(in_channels)
  
      def forward(self, x):
  
          x = self.depthwise(x)
          x = self.relu(x)
          x = self.bn(x)
          return x
  
  class PointWiseSeperable(nn.Module):
  
      def __init__(self, in_channels, out_channels, stride):
  
          super(PointWiseSeperable, self).__init__()
  
          self.pointwise =  nn.Conv2d(in_channels = in_channels , out_channels = out_channels , stride = 1 , padding = 0, kernel_size = 1, bias = False)
          self.relu = nn.PReLU(out_channels)
          self.bn = nn.BatchNorm2d(out_channels)
  
      def forward(self, x):
  
          x = self.pointwise(x)
          x = self.relu(x)
          x = self.bn(x)
          return x
  
  class SepConv2d(nn.Module):
      def __init__(self, in_channels, out_channels, stride, padding, kernel_size):
  
          super(SepConv2d, self).__init__()
  
          self.depthwise = DepthWiseSeperable(in_channels, in_channels, stride, padding, kernel_size)
          self.pointwise = PointWiseSeperable(in_channels, out_channels, stride)
  
      def forward(self, x):
          x = self.depthwise(x)
          x = self.pointwise(x)
  
          return x
  
  ```
  
  ---
  
  Since later on we are going to create an instance of MobileNetV1 from huggingface a class is implemented to truncate such net. Paper says:
  
  > layers after the depthwise convolution of block 9 are skipped
  
  ```
  class Truncated_MobileNet(nn.Module):
      
      def __init__(self, original_model):
          super(Truncated_MobileNet, self).__init__()
          self.conv_stem = original_model.conv_stem
          self.layer = nn.ModuleList(original_model.layer[:17])  #layers after the depthwise convolution of block 9 are skipped
  
      def forward(self, x):
          x = self.conv_stem(x)
          for layer in self.layer:
              x = layer(x)
    
          return x
  
  ```
  
  ---
  
  Here patch extraction module. In the paper they say that pading is needed from 14x14 to 16x16, but I already obtain 16x16 from truncated mobilenet. I adjust the kernel sizes and strides to match dimension with the architecture graphic.
  
  ```
  class PatchExtraction(nn.Module):
  
      def __init__(self, in_channels, out_channels):
          """
          Patch extraction block.
  
          Args:
              in_channels (int): Number of input channels.
              out_channels (int): Number of output channels.
          """
          super(PatchExtraction, self).__init__()
  
          # Padding to achieve 16x16 spatial dimensions
          self.padding = nn.ZeroPad2d((1, 1, 1, 1))
  
          # Depthwise separable convolution output 4x4x256
          self.depthwise1 = SepConv2d(in_channels = 512, out_channels = 256, stride = 4 , padding = 0, kernel_size = 4)
          # Depthwise separable convolution output 2x2x256
          self.depthwise2 = SepConv2d(in_channels = 256, out_channels = 256, stride = 2 , padding = 0, kernel_size = 2)
  
          # Pointwise convolution output 2x2x256
          self.pointwise = PointWiseSeperable(in_channels = 256, out_channels = 256, stride=1)
  
      def forward(self, x):
  
          x = self.depthwise1(x)
          x = self.depthwise2(x)
          x = self.pointwise(x)
  
          return x
  
  ```
  
  ---
  
  The attention classifier module. **I am using residual as I saw that stabilises the training, but I don't know if this way is the best practice.**
  
  ```
  class AttentionClassifier(nn.Module):
      def __init__(self, in_features, hidden_size, num_classes, d_q):
          super(AttentionClassifier, self).__init__()
          
          # First fully connected layer
          self.fc1 = nn.Linear(in_features, hidden_size)
          self.relu = nn.ReLU()
  
          # Linear projections for query, key, and value
          self.fc_query = nn.Linear(hidden_size, d_q)
          self.fc_key = nn.Linear(hidden_size, d_q)
          self.fc_value = nn.Linear(hidden_size, d_q)  
  
      
          # Second fully connected layer
          self.fc2 = nn.Linear(d_q, num_classes)
          #torch.nn.init.xavier_uniform_(self.fc2.weight)
          self.softmax = nn.Softmax(dim=-1)
          # ScaledDotProductAttention instance
          #self.attention = ScaledDotProductAttention(d_q)
  
      def forward(self, x):
          
          # Apply the first fully connected layer
          y = self.fc1(x)
          y = self.relu(y)
          # Linear projections for query, key, and value
          query = self.fc_query(y)
          key = self.fc_key(y)
          value = self.fc_value(y)
  
          attention_output = F.scaled_dot_product_attention(query, key, value, dropout_p= 0.0, scale = None)
          attention_output = attention_output + x
          # Apply the second fully connected layer
          x = self.softmax(self.fc2(attention_output))
          
          return x
  
  ```
  
  ---
  
  Joining patch extraction and attention modules. **I don't know if input\_dim<hidden\_dim is the best use of attention.**
  
  ```
  class Top_model(nn.Module):
      def __init__(self,  num_classes=7):
          
          super(Top_model, self).__init__()
          
          self.patchextraction = PatchExtraction(512, 256)
          
          # Global Average Pooling and classifier
  
          self.classifier = nn.Sequential(
              AttentionClassifier(in_features=256, hidden_size=512, num_classes=7, d_q=256),
          )
          self.gap = nn.AdaptiveAvgPool2d((1,1))
  
          
      def forward(self, x):
          x = self.patchextraction(x)
          x = self.gap(x)
          x = x.view(x.size(0), -1)  # Flattening the tensor
          x = self.classifier(x)
  
          return x    
  
  ```
  
  ---
  
  Finally join all:
  
  ```
  class Patt_lite(nn.Module):
      
      def __init__(self, truncated_model, top_model, freeze_base=True, freeze_top = True):
          
          super(Patt_lite, self).__init__()
  
          self.base = truncated_model
          if freeze_base:
              for i, param in enumerate(self.base.parameters()):
                  param.requires_grad = False
          
          self.top = top_model
          if freeze_top:
              for i, param in enumerate(self.top.parameters()):
                  param.requires_grad = False
          
      def forward(self, x):
          x = self.base(x)
          x = self.top(x)
  
          return x      
  
  ```
  
  ---
  
  For the training I'm using Fastai:
  
  ```
  # Data augmentation for training
  train_batch_transforms = [
       # Quality
       Brightness(p=0), Contrast(p=0), Saturation(p=0, draw = 0),
       # Data Alterations
       RandomErasing(p=0.25),
       # Orientation
       Flip(p=0.5), Rotate(p=0.75, draw=30), Warp(p=0.5),
       #Norm
       Normalize.from_stats([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
   ]
  train_item_transforms = [
      Resize(224),
      ToTensor(),
  ]
  
  test_batch_transforms = [
       Normalize.from_stats([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
   ]
  test_item_transforms = [
      #ToTensor, Normalize
      Resize(224),
      ToTensor(),
      
  ]
  
  #Set batch size
  batch_size = 8
  
  # Create data loaders using the defined transforms
  dls = ImageDataLoaders.from_folder(
      path='path to train dataset',
      valid_pct=0.2,
      bs=batch_size,
      valid_bs=batch_size,
      item_tfms=train_item_transforms,
      batch_tfms=train_batch_transforms,
      device="cuda",
      convert_mode='L'
  )
  
  #Model instance
  configuration = MobileNetV1Config(tf_padding=False)
  original_model = MobileNetV1Model.from_pretrained("Matthijs/mobilenet_v1_1.0_224", config=configuration)
  truncated_model = Truncated_MobileNet(original_model)
  
  top_model = Top_model(num_classes = 7)
  patt_model = Patt_lite(truncated_model, top_model, freeze_base = False, freeze_top = False)
  patt_model = patt_model.to(device)
  
  #Optimization and learning
  opt_func = partial(OptimWrapper, opt=optim.RMSprop)
  learn = Learner(dls, patt_model, opt_func=opt_func, metrics=accuracy)
  learn.loss_func = CrossEntropyLossFlat()
  learn.fit(n_epoch=50, lr = 1e-4, cbs = [ShowGraphCallback()])
  
  ```
  
  I tried `Matthijs/mobilenet_v1_1.0_224` and `google/mobilenet_v1_1.0_224`, but seems the same.
  
  On the paper they split the training in two parts: top model training with mobilenet freezed and fine tuning of mobilenet. **Are the weights of mobilenet correctly loaded?**
  
  The first training result in bad performance and the second one do not improve the accuracy or loss.
  
  [![enter image description here](https://i.sstatic.net/CzdrC.png)](https://i.sstatic.net/CzdrC.png)
  [![enter image description here](https://i.sstatic.net/hqTz8.png)](https://i.sstatic.net/hqTz8.png)
  
  Any idea? Thanks
  
  ---
  
  Edit:
  Option A: Eliminate dot\_product\_attention (also softmax and residual), so
  
  ```
  class AttentionClassifier(nn.Module):
      def __init__(self, in_features, hidden_size, num_classes, d_q):
          super(AttentionClassifier, self).__init__()
          
          # First fully connected layer
          self.fc1 = nn.Linear(in_features, hidden_size)
          self.relu = nn.ReLU()
  
     
      
          # Second fully connected layer
          self.fc2 = nn.Linear(d_q, num_classes)
  
      def forward(self, x):
          
          # Apply the first fully connected layer
          x = self.fc1(x)
          x = self.relu(x)
         
          # Apply the second fully connected layer
          x = self.fc2(x)
          #x = self.softmax(x)
          
          return x
  
  ```
  
  After 25 epoch:
  [![enter image description here](https://i.sstatic.net/cLhlD.png)](https://i.sstatic.net/cLhlD.png)
  [![enter image description here](https://i.sstatic.net/ywW93.png)](https://i.sstatic.net/ywW93.png)
  
  ---
  
  Option B: No patch, no dot\_product\_attention, just MobileNet (freezed) and 2 FC layers.
  [![enter image description here](https://i.sstatic.net/Qd0ke.png)](https://i.sstatic.net/Qd0ke.png)
  [![enter image description here](https://i.sstatic.net/WMtc1.png)](https://i.sstatic.net/WMtc1.png)
  
  ---
  
  Option C: same as Option B but with mobilenet unfreezed. Clearly overfitting.
  [![enter image description here](https://i.sstatic.net/zbUYo.png)](https://i.sstatic.net/zbUYo.png)
  [![enter image description here](https://i.sstatic.net/zll0R.png)](https://i.sstatic.net/zll0R.png)
  
  ---
  
  Option D: No patch extraction, attention without residual.
  
  ```
  class AttentionClassifier(nn.Module):
      def __init__(self, in_features, hidden_size, num_classes, d_q):
          super(AttentionClassifier, self).__init__()
          
          # First fully connected layer
          self.fc1 = nn.Linear(in_features, hidden_size)
          self.relu = nn.ReLU()
  
          # Linear projections for query, key, and value
          self.fc_query = nn.Linear(hidden_size, d_q)
          self.fc_key = nn.Linear(hidden_size, d_q)
          self.fc_value = nn.Linear(hidden_size, d_q)  
  
      
          # Second fully connected layer
          self.fc2 = nn.Linear(d_q, num_classes)
  
      def forward(self, x):
          
          # Apply the first fully connected laye
          x = self.fc1(x)
          x = self.relu(x)
          # Linear projections for query, key, and value
          query = self.fc_query(x)
          key = self.fc_key(x)
          value = self.fc_value(x)
  
          attention_output = F.scaled_dot_product_attention(query, key, value, dropout_p= 0.0, scale = None)
          #Apply the second fully connected layer
          x = self.fc2(attention_output)
          #x = self.softmax(x)
          
          return x
  
      
  class Top_model(nn.Module):
      
      def __init__(self,  num_classes=7):
          
          super(Top_model, self).__init__()
          
          self.patchextraction = PatchExtraction(512, 256)
          
          # Global Average Pooling and classifier
  
          self.classifier = nn.Sequential(
              AttentionClassifier(in_features=512, hidden_size=256, num_classes=num_classes, d_q=256),
          )
          self.gap = nn.AdaptiveAvgPool2d((1,1))
  
          
      def forward(self, x):
          x = self.gap(x)
          x = x.view(x.size(0), -1)  # Flattening the tensor
          x = self.classifier(x)
  
          return x   
  
  ```
  
  [![enter image description here](https://i.sstatic.net/11jFG.png)](https://i.sstatic.net/11jFG.png)
  [![enter image description here](https://i.sstatic.net/5zOPL.png)](https://i.sstatic.net/5zOPL.png)
  
  ---
  
  Option E: Same as option D but with  `attention_output = attention_output + x`
  
  [![enter image description here](https://i.sstatic.net/scc8F.png)](https://i.sstatic.net/scc8F.png)
  [![enter image description here](https://i.sstatic.net/n0j5F.png)](https://i.sstatic.net/n0j5F.png)
  
  **Tags:** computer-vision, attention, transfer-learning, facial-recognition, state-of-the-art
  
  > **Comments:**
  > **username_1**: In my opinion it's always important to make sure that the model trains first.   Since your training loss does not really decrease, it's interesting to know if you have tried changing the hyperparameters? For example increasing the learning rate?   I also noticed that in your attention layer you are adding x as a residual element and not the y, that looks like an input to attention head. Is it intentional?  
  > **username_1**: I don't see any problems in the code at first glance, but what can be useful, if you have in fact checked that its not hyperparametetrs problem, is to try simpler base or top model. For top you can just flatten all the features and apply one fully connected layer for example. Like this you can assure yourself that specific parts of your model in fact are working as they are suppose to. You can really do this for all the important "elements" of the model, where the problem can be hidden in your opinion.  
  > **username_0**: The using of x and y was intentional since i tried to increase the hidden dimension, so the hidden dimension was grater than k,q, v and attention output dimension, thus y = self.fc1(x) does not match with attention_output but x matches. I tried also with input_dim = hidden_dim = output_dim and use y as a residual: same performance.  
  > **username_0**: Thanks! I will try it part by part to address the issue.  
  > **username_0**: @username_1 Do i need a flatten between GlobalAvgPool and the classifier aprt in top model forward? It is enough with the AdaptativeAvgPool2d((1,1)) and x = x.view(x.size(0), -1)??  
  > **username_1**: I would suggest you replace x.view(x.size(0), -1) by torch.flatten (https://pytorch.org/docs/stable/generated/torch.flatten.html)  
  > **username_1**: Can you try doing global average polling on feature maps of mobile net and training your attention classifier first, to see if it works? Without the "patch extraction block"?  
  > **username_1**: Let us [continue this discussion in chat](https://chat.stackexchange.com/rooms/151468/discussion-between-username_1-and-username_0).  
  
  ---
  
  ## 0 answers
  
  