File size: 6,952 Bytes
e300024 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | """Native PyTorch implementation of the Perch v2 EfficientNet-B3 model.
Note: this file has been entirely generated by chatGPT 5.6 with codex, and I did not review every line.
I take the entire responsability for any error.
This file is the pytorch translation of the onnx artefact published here: https://www.kaggle.com/datasets/nikitababich/perchv2-onnx
You can download it locally with:
```
path = kagglehub.dataset_download("nikitababich/perchv2-onnx")
```
To generate the `.pth` weight file from the onnx, use `tools/extract_perch_backbone.py`
"""
from __future__ import annotations
import math
import torch
import torch.nn.functional as F
from torch import Tensor, nn
_WIDTH = 1.2
_DEPTH = 1.4
_NUM_PERCH_CLASSES = 14795
_NUM_PROTOTYPES = 4
_STAGES = (
(1, 16, 3, 1, 1),
(2, 24, 3, 2, 6),
(2, 40, 5, 2, 6),
(3, 80, 3, 2, 6),
(3, 112, 5, 1, 6),
(4, 192, 5, 2, 6),
(1, 320, 3, 1, 6),
)
def _channels(channels: int) -> int:
scaled = channels * _WIDTH
rounded = max(8, int(scaled + 4) // 8 * 8)
return int(rounded + 8 if rounded < 0.9 * scaled else rounded)
def _blocks(count: int) -> int:
return math.ceil(count * _DEPTH)
def _same_padding(x: Tensor, kernel_size: int, stride: int) -> Tensor:
height, width = x.shape[-2:]
def padding(size: int) -> tuple[int, int]:
total = max((math.ceil(size / stride) - 1) * stride + kernel_size - size, 0)
return total // 2, total - total // 2
top, bottom = padding(height)
left, right = padding(width)
return F.pad(x, (left, right, top, bottom))
torch.fx.wrap("_same_padding")
class SqueezeExcitation(nn.Module):
"""Per-channel squeeze-and-excitation gate."""
def __init__(self, channels: int, reduced_channels: int) -> None:
super().__init__()
self.reduce = nn.Linear(channels, reduced_channels)
self.expand = nn.Linear(reduced_channels, channels)
def forward(self, x: Tensor) -> Tensor:
scale = torch.sigmoid(self.expand(F.silu(self.reduce(x.mean(dim=(-2, -1))))))
return x * scale[:, :, None, None]
class MBConv(nn.Module):
"""Mobile inverted bottleneck block used by Perch."""
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int, expansion: int) -> None:
super().__init__()
expanded_channels = in_channels * expansion
self.kernel_size = kernel_size
self.stride = stride
self.has_expand = expansion != 1
if self.has_expand:
self.expand_conv = nn.Conv2d(in_channels, expanded_channels, 1, bias=False)
self.expand_bn = nn.BatchNorm2d(expanded_channels)
self.depthwise_conv = nn.Conv2d(
expanded_channels,
expanded_channels,
kernel_size,
stride=stride,
groups=expanded_channels,
bias=False,
)
self.depthwise_bn = nn.BatchNorm2d(expanded_channels)
self.se = SqueezeExcitation(expanded_channels, expanded_channels // (4 * expansion))
self.project_conv = nn.Conv2d(expanded_channels, out_channels, 1, bias=False)
self.project_bn = nn.BatchNorm2d(out_channels)
def forward(self, x: Tensor) -> Tensor:
if self.has_expand:
x = F.silu(self.expand_bn(self.expand_conv(x)))
x = _same_padding(x, self.kernel_size, self.stride)
x = F.silu(self.depthwise_bn(self.depthwise_conv(x)))
return self.project_bn(self.project_conv(self.se(x)))
class ResidualMBConv(nn.Module):
"""MBConv with its eligible residual connection."""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int,
expansion: int,
residual: bool,
) -> None:
super().__init__()
self.mbconv = MBConv(in_channels, out_channels, kernel_size, stride, expansion)
self.residual = residual
def forward(self, x: Tensor) -> Tensor:
result = self.mbconv(x)
return result + x if self.residual else result
class PerchBackbone(nn.Module):
"""Perch v2 feature extractor with selectable ``blocks.<index>`` layers."""
def __init__(self, input_layout: str = "frequency_time") -> None:
super().__init__()
if input_layout not in {"frequency_time", "time_frequency"}:
raise ValueError("input_layout must be 'frequency_time' or 'time_frequency'")
self.input_layout = input_layout
stem_channels = _channels(32)
self.stem_conv = nn.Conv2d(1, stem_channels, 3, stride=2, bias=False)
self.stem_bn = nn.BatchNorm2d(stem_channels)
blocks: list[nn.Module] = []
in_channels = stem_channels
for stage_blocks, stage_channels, kernel_size, stage_stride, expansion in _STAGES:
out_channels = _channels(stage_channels)
for block_index in range(_blocks(stage_blocks)):
blocks.append(
ResidualMBConv(
in_channels,
out_channels,
kernel_size,
stage_stride if block_index == 0 else 1,
expansion,
residual=block_index > 0,
)
)
in_channels = out_channels
self.blocks = nn.ModuleList(blocks)
self.head_conv = nn.Conv2d(in_channels, _channels(1280), 1, bias=False)
self.head_bn = nn.BatchNorm2d(_channels(1280))
def forward(self, x: Tensor) -> Tensor:
if self.input_layout == "time_frequency":
x = x.transpose(-1, -2).contiguous()
x = F.silu(self.stem_bn(self.stem_conv(x)))
for block in self.blocks:
x = block(x)
return F.silu(self.head_bn(self.head_conv(x)))
class ProtoPNetHead(nn.Module):
"""Perch v2 ProtoPNet classifier."""
def __init__(self) -> None:
super().__init__()
self.prototypes = nn.Parameter(torch.empty(_NUM_PERCH_CLASSES, _channels(1280), _NUM_PROTOTYPES))
self.kernel = nn.Parameter(torch.empty(_NUM_PERCH_CLASSES, _NUM_PROTOTYPES))
self.bias = nn.Parameter(torch.empty(_NUM_PERCH_CLASSES))
def forward(self, spatial_embedding: Tensor) -> Tensor:
normalized_embedding = spatial_embedding / (spatial_embedding.norm(dim=1, keepdim=True) + 1e-5)
similarities = torch.einsum("bdhw,cdp->bhwcp", normalized_embedding, self.prototypes).amax(dim=(1, 2))
return (similarities * self.kernel.clamp_min(0)).sum(dim=-1) + self.bias
class PerchModel(nn.Module):
"""Perch v2 backbone and ProtoPNet classifier."""
def __init__(self, input_layout: str = "frequency_time") -> None:
super().__init__()
self.backbone = PerchBackbone(input_layout)
self.classifier = ProtoPNetHead()
def forward(self, x: Tensor) -> Tensor:
return self.classifier(self.backbone(x))
|