1
$\begingroup$

I'm trying to implement a model guided by the paper of PAtt-lite (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

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 enter image description here

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 enter image description here


Option B: No patch, no dot_product_attention, just MobileNet (freezed) and 2 FC layers. enter image description here enter image description here


Option C: same as Option B but with mobilenet unfreezed. Clearly overfitting. enter image description here enter image description here


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 enter image description here


Option E: Same as option D but with attention_output = attention_output + x

enter image description here enter image description here

$\endgroup$
8
  • $\begingroup$ 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? $\endgroup$
    – vl_knd
    Commented Feb 13 at 13:45
  • $\begingroup$ 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. $\endgroup$
    – vl_knd
    Commented Feb 13 at 13:48
  • $\begingroup$ 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. $\endgroup$
    – Robert
    Commented Feb 13 at 13:58
  • $\begingroup$ Thanks! I will try it part by part to address the issue. $\endgroup$
    – Robert
    Commented Feb 13 at 14:00
  • $\begingroup$ @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)?? $\endgroup$
    – Robert
    Commented Feb 13 at 14:03

0