Pytorch - timm, torchvision.models part2)

In [1]:

from IPython.core.display import display, HTML

display(HTML("<style>.container { width:140% !important; }</style>"))

View Source

PyTorch Model - timm library, torchvision.models part2)¶

Application of timm and torchvision.model libraries

In [2]:

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from torchvision import models

Timm Library¶

Creating Custom Model using convnext_small¶

Implementing Custom Model for ThreeWayStratifiedDataset described in [https://psyduck5.tistory.com/10]
Attempt1. Sharing pretrained embeddings but constructing differnet layers for each labels (age, mask, gender)
Attempt2. For Generalization, applying dropout layers
Goal. Not Enough GPU resources to train each labels. Therefore, by implementing Attempt1, Attemp2, trying the acheive the ensemble-like effects

In [3]:

class ConvNextLSmallCustom(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = timm.create_model('convnext_small', pretrained=True, num_classes = 1536)
        self.dropout = nn.Dropout(0.5)
        self.dropouts = nn.ModuleList([
                    nn.Dropout(0.5) for _ in range(5)])
        self.age_layer = nn.Linear(in_features=1536, out_features=3, bias=True)
        self.mask_layer = nn.Linear(in_features=1536, out_features=2, bias=True)
        self.sex_layer = nn.Linear(in_features=1536, out_features=2, bias=True)
    
    def forward(self, x):
        x = self.model(x)
        x_ = self.dropout(x)
        
        for i, dropout in enumerate(self.dropouts):
            if i==0:
                x_age = self.age_layer(dropout(x_))
                x_mask = self.mask_layer(dropout(x_))
                x_sex = self.sex_layer(dropout(x_))
            else:
                x_age += self.age_layer(dropout(x_))
                x_mask += self.mask_layer(dropout(x_))
                x_sex += self.sex_layer(dropout(x_))
        else:
            x_age /= len(self.dropouts)
            x_mask /= len(self.dropouts)
            x_sex /= len(self.dropouts)
        
        return x_age, x_mask, x_sex

In [9]:

model = ConvNextLSmallCustom()
model.eval()
result = model(torch.randn(1,3,224,224))
print(result)

(tensor([[0.0389, 0.0142, 0.0060]], grad_fn=<DivBackward0>), tensor([[ 0.0206, -0.0145]], grad_fn=<DivBackward0>), tensor([[-0.0092, -0.0465]], grad_fn=<DivBackward0>))

Torchvision models¶

Initialzing with functions¶

Implementing user-defined functions for changing the last layer and initializing weights
Attempt1. Define a function to automatically finds out the last layer of the model and change the out_features
Attempt2. Define a function to initialzing wegihts to the modified layers.
Goal. With these functions, aimimg to contruct the custom model class just like timm library

In [4]:

def initialize_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        nn.init.zeros_(m.bias)

    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_uniform_(m.weight.data, nonlinearity='relu') 
    
def change_last_layer(model, num_classes):
    name_last_layer = list(model.named_modules())[-1][0]
    
    if name_last_layer == 'classifier':
        model.classifier = nn.Linear(in_features = model.classifier.in_features,
                                              out_features = num_classes, bias = True)
        initialize_weights(model.classifier)
        return model
    
    elif name_last_layer == 'fc':
        model.fc = nn.Linear(in_features = model.fc.in_features,
                                      out_features = num_classes, bias = True)
        initialize_weights(model.fc)
        return model
    
    else:
        raise Exceptionception('last layer should be nn.Linear Module named as either fc or classifier')

In [5]:

class DenseNet161(nn.Module):
    def __init__(self, num_classes):
        super(DenseNet161, self).__init__()
        self.model = models.densenet161(pretrained = True)
        self.model = change_last_layer(self.model, num_classes)
    
    def forward(self, x):
        x = self.model(x)
        return x

In [6]:

model = DenseNet161(10)
model.eval()
model(torch.randn(1,3,224,224))

Out[6]:

tensor([[ 0.2204,  0.6942, -0.8310,  0.0541, -1.7376, -0.1588,  0.0203,  0.7483,
          1.5705, -0.7220]], grad_fn=<AddmmBackward0>)