| """Algebraic pointwise-Conv -> Linear lowering; unchanged masked SE semantics.""" |
| import copy |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| def linear(conv): |
| layer = nn.Linear(conv.in_channels, conv.out_channels) |
| with torch.no_grad(): |
| layer.weight.copy_(conv.weight[..., 0]) |
| layer.bias.copy_(conv.bias) |
| return layer |
|
|
| class PortableBlock(nn.Module): |
| def __init__(self, block): |
| super().__init__() |
| self.dw = copy.deepcopy(block.dw) |
| for name in ['expand', 'fc1', 'fc2', 'project']: |
| setattr(self, name, linear(getattr(block, name))) |
|
|
| def forward(self, x, mask): |
| z = F.hardswish(self.expand(self.dw(x.transpose(1,2)).transpose(1,2))) |
| pool = (z * mask).sum(1, keepdim=True) / mask.sum(1, keepdim=True).clamp_min(1.) |
| gate = torch.sigmoid(self.fc2(F.relu(self.fc1(pool)))) |
| return (x + self.project(z * gate)) * mask |
|
|
| class PortableRecognizer(nn.Module): |
| """Input widths are true unpadded widths; canvas must be a multiple of eight.""" |
| def __init__(self, source): |
| super().__init__() |
| self.index_dtype = torch.int64 |
| self.stem = copy.deepcopy(source.stem) |
| self.bn = copy.deepcopy(source.bn) |
| self.blocks = nn.ModuleList(PortableBlock(b) for b in source.blocks) |
| self.head = linear(source.head) |
|
|
| def forward(self, pixels, widths): |
| z = self.stem(F.pad(pixels, (0,4,0,0))).squeeze(2) |
| lengths = (widths + 3) // 4 |
| mask = (torch.arange(z.shape[-1], device=z.device, dtype=self.index_dtype)[None,:] < lengths[:,None])[:,:,None].to(z.dtype) |
| z = F.hardswish(self.bn(z)).transpose(1,2) * mask |
| for block in self.blocks: |
| z = block(z, mask) |
| return self.head(z), lengths |
|
|
| class Compact(nn.Module): |
| def __init__(self, source): |
| super().__init__() |
| self.model = source |
| def forward(self, pixels, widths): |
| logits, lengths = self.model(pixels, widths) |
| confidence, tokens = logits.softmax(-1).max(-1) |
| return tokens.to(torch.int32), confidence, lengths.to(torch.int32) |
|
|
