Create model.py

tasks/model.py

@@ -0,0 +1,107 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from torchaudio import transforms
+from torchvision import models
+
+
+
+class BlazeFace(nn.Module):
+    def __init__(self, input_channels=1, use_double_block=False, activation="relu", use_optional_block=True):
+        super(BlazeFace, self).__init__()
+        self.activation = activation
+        self.use_double_block = use_double_block
+        self.use_optional_block = use_optional_block
+
+        def conv_block(in_channels, out_channels, kernel_size, stride, padding):
+            return nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU() if activation == "relu" else nn.Sigmoid()  # Apply ReLU activation (default) or Sigmoid
+            )
+        
+        def depthwise_separable_block(in_channels, out_channels, stride):
+            return nn.Sequential(
+                nn.Conv2d(in_channels, in_channels, kernel_size=5, stride=stride, padding=2, groups=in_channels, bias=False),
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU() if activation == "relu" else nn.Sigmoid()
+            )
+
+        def double_block(in_channels, filters_1, filters_2, stride):
+            return nn.Sequential(
+                depthwise_separable_block(in_channels, filters_1, stride),
+                depthwise_separable_block(filters_1, filters_2, 1)
+            )
+
+        # Define layers (first part: conv layers)
+        self.conv1 = conv_block(input_channels, 24, kernel_size=5, stride=2, padding=2)
+
+        # Define single blocks (subsequent conv blocks)
+        self.single_blocks = nn.ModuleList([
+            depthwise_separable_block(24, 24, stride=1),
+            depthwise_separable_block(24, 24, stride=1),
+            depthwise_separable_block(24, 48, stride=2),
+            depthwise_separable_block(48, 48, stride=1),
+            depthwise_separable_block(48, 48, stride=1)
+        ])
+
+        # Define double blocks if `use_double_block` is True
+        if self.use_double_block:
+            self.double_blocks = nn.ModuleList([
+                double_block(48, 24, 96, stride=2),
+                double_block(96, 24, 96, stride=1),
+                double_block(96, 24, 96, stride=2),
+                double_block(96, 24, 96, stride=1),
+                double_block(96, 24, 96, stride=2)
+            ])
+        else:
+            self.double_blocks = nn.ModuleList([
+                depthwise_separable_block(48, 96, stride=2),
+                depthwise_separable_block(96, 96, stride=1),
+                depthwise_separable_block(96, 96, stride=2),
+                depthwise_separable_block(96, 96, stride=1),
+                depthwise_separable_block(96, 96, stride=2)
+            ])
+        
+        # Final convolutional head
+        self.conv_head = nn.Conv2d(96, 64, kernel_size=1, stride=1)
+        self.bn_head = nn.BatchNorm2d(64)
+
+        # Global Average Pooling
+        self.global_avg_pooling = nn.AdaptiveAvgPool2d(1)
+
+    def forward(self, x):
+        # First conv layer
+        x = self.conv1(x)
+
+        # Apply single blocks
+        for block in self.single_blocks:
+            x = block(x)
+
+        # Apply double blocks
+        for block in self.double_blocks:
+            x = block(x)
+
+        # Final head
+        x = self.conv_head(x)
+        x = self.bn_head(x)
+        x = F.relu(x)
+
+        # Global Average Pooling and Flatten
+        x = self.global_avg_pooling(x)
+        x = torch.flatten(x, 1)
+
+        return x
+
+class BlazeFaceModel(nn.Module):
+    def __init__(self, input_channels, label_count, use_double_block=False, activation="relu", use_optional_block=True):
+        super(BlazeFaceModel, self).__init__()
+        self.blazeface_backbone = BlazeFace(input_channels=input_channels, use_double_block=use_double_block, activation=activation, use_optional_block=use_optional_block)
+        self.fc = nn.Linear(64, label_count)  
+
+    def forward(self, x):
+        features = self.blazeface_backbone(x)
+        output = self.fc(features)
+        return output