Spaces:
Sleeping
Sleeping
| 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 | |
| from torch.quantization import QuantStub, DeQuantStub | |
| 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 | |
| class QuantizedBlazeFaceModel(nn.Module): | |
| def __init__(self, model_fp32): | |
| super(QuantizedBlazeFaceModel, self).__init__() | |
| self.quant = QuantStub() | |
| self.dequant = DeQuantStub() | |
| self.backbone = model_fp32.blazeface_backbone | |
| self.fc = model_fp32.fc | |
| def forward(self, x): | |
| x = self.quant(x) | |
| x = self.backbone(x) | |
| x = self.fc(x) | |
| x = self.dequant(x) | |
| return x |