| """Width512/depth20 Indra-family CTC recognizer; native class IDs are retained.""" |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.utils.checkpoint import checkpoint |
| class Block(nn.Module): |
| def __init__(self,width=512,identity=False): |
| super().__init__();self.dw=nn.Conv1d(width,width,17,padding=8,groups=width);self.expand=nn.Conv1d(width,2*width,1);self.fc1=nn.Conv1d(2*width,width//2,1);self.fc2=nn.Conv1d(width//2,2*width,1);self.project=nn.Conv1d(2*width,width,1) |
| if identity:nn.init.zeros_(self.project.weight);nn.init.zeros_(self.project.bias) |
| def forward(self,x,mask): |
| z=F.hardswish(self.expand(self.dw(x)));pool=(z*mask).sum(2,keepdim=True)/mask.sum(2,keepdim=True).clamp_min(1);return(x+self.project(z*torch.sigmoid(self.fc2(F.relu(self.fc1(pool))))))*mask |
| class Recognizer(nn.Module): |
| input_height=64 |
| def __init__(self,classes=1357,width=512,depth=20,gradient_checkpointing=True): |
| super().__init__();assert width==512 and depth==20 |
| self.stem=nn.Conv2d(1,width,(64,8),stride=(64,4));self.bn=nn.BatchNorm1d(width);self.blocks=nn.ModuleList(Block(width,identity=i%2==1)for i in range(depth));self.head=nn.Conv1d(width,classes,1);self.gradient_checkpointing=gradient_checkpointing;self.bn.eval() |
| def train(self,mode=True):super().train(mode);self.bn.eval();return self |
| def forward(self,x,widths): |
| if x.ndim!=4 or x.shape[1]!=1 or x.shape[2]!=64:raise ValueError('Require Bx1x64xW without height truncation') |
| if widths.ndim!=1 or len(widths)!=len(x)or bool((widths<1).any())or bool((widths>x.shape[-1]).any()):raise ValueError('Invalid original crop widths') |
| widths=widths.to(device=x.device,dtype=torch.long) |
| if x.shape[-1]%4:x=F.pad(x,(0,(-x.shape[-1])%4)) |
| z=self.stem(F.pad(x,(0,4,0,0))).squeeze(2);lengths=(widths+3)//4;assert bool((lengths<=z.shape[-1]).all());mask=(torch.arange(z.shape[-1],device=z.device)[None,:]<lengths[:,None])[:,None,:].to(z.dtype);z=F.hardswish(self.bn(z))*mask |
| for block in self.blocks: |
| z=checkpoint(block,z,mask,use_reentrant=False)if self.training and self.gradient_checkpointing and torch.is_grad_enabled()else block(z,mask) |
| return self.head(z).transpose(1,2),lengths |
| def lift_from_highres(self,source): |
| """Net2Wider copies old128 features fourfold; new interleaved blocks start as identities.""" |
| state=source.state_dict()if isinstance(source,nn.Module)else source |
| assert state['stem.weight'].shape==(128,1,64,8)and state['blocks.9.dw.weight'].shape==(128,1,17) |
| with torch.no_grad(): |
| self.stem.weight.copy_(state['stem.weight'].repeat(4,1,1,1));self.stem.bias.copy_(state['stem.bias'].repeat(4)) |
| for name in ['weight','bias','running_mean','running_var']:getattr(self.bn,name).copy_(state['bn.'+name].repeat(4)) |
| self.bn.num_batches_tracked.copy_(state['bn.num_batches_tracked']) |
| for i in range(10): |
| b=self.blocks[2*i];prefix='blocks.'+str(i)+'.';b.dw.weight.copy_(state[prefix+'dw.weight'].repeat(4,1,1));b.dw.bias.copy_(state[prefix+'dw.bias'].repeat(4)) |
| for name in ['expand','fc1','fc2','project']: |
| layer=getattr(b,name);layer.weight.copy_(state[prefix+name+'.weight'].repeat(4,4,1)/4);layer.bias.copy_(state[prefix+name+'.bias'].repeat(4)) |
| n=state['head.weight'].shape[0];assert n<=self.head.out_channels |
| self.head.weight[:n].copy_(state['head.weight'].repeat(1,4,1)/4);self.head.bias[:n].copy_(state['head.bias']) |
| if n<self.head.out_channels:self.head.weight[n:].zero_();self.head.bias[n:].fill_(-8.) |
|
|
