| import torch, torch.nn as nn |
| class DW(nn.Module): |
| def __init__(s,i,o,st,exp=2): |
| super().__init__(); h=i*exp; s.res=(st==1 and i==o) |
| s.b=nn.Sequential(nn.Conv2d(i,h,1,bias=False),nn.BatchNorm2d(h),nn.PReLU(h), |
| nn.Conv2d(h,h,3,st,1,groups=h,bias=False),nn.BatchNorm2d(h),nn.PReLU(h), |
| nn.Conv2d(h,o,1,bias=False),nn.BatchNorm2d(o)) |
| def forward(s,x): return x+s.b(x) if s.res else s.b(x) |
| class TinyFaceNet(nn.Module): |
| """~1.0M params, 112x112 in, 128-d embedding.""" |
| def __init__(s,emb=128,n_cls=10,w=1.6): |
| super().__init__(); c=lambda v:int(v*w) |
| s.stem=nn.Sequential(nn.Conv2d(3,c(32),3,2,1,bias=False),nn.BatchNorm2d(c(32)),nn.PReLU(c(32)), |
| nn.Conv2d(c(32),c(32),3,1,1,groups=c(32),bias=False),nn.BatchNorm2d(c(32)),nn.PReLU(c(32))) |
| cfg=[(c(32),c(48),2),(c(48),c(48),1),(c(48),c(48),1),(c(48),c(96),2),(c(96),c(96),1),(c(96),c(96),1),(c(96),c(96),1), |
| (c(96),c(128),2),(c(128),c(128),1),(c(128),c(128),1)] |
| s.blocks=nn.Sequential(*[DW(i,o,st) for i,o,st in cfg]) |
| s.head=nn.Sequential(nn.Conv2d(c(128),c(256),1,bias=False),nn.BatchNorm2d(c(256)),nn.PReLU(c(256)), |
| nn.Conv2d(c(256),c(256),7,1,0,groups=c(256),bias=False),nn.BatchNorm2d(c(256)), |
| nn.Flatten(),nn.Linear(c(256),emb,bias=False),nn.BatchNorm1d(emb)) |
| s.distill=nn.Linear(emb,512) |
| s.cls=nn.Linear(emb,n_cls) |
| def forward(s,x): return s.head(s.blocks(s.stem(x))) |
| if __name__=="__main__": |
| m=TinyFaceNet(); core=sum(p.numel() for n,p in m.named_parameters() if not n.startswith(("distill","cls"))) |
| print("deploy params",core, "total",sum(p.numel() for p in m.parameters())); print(m(torch.zeros(2,3,112,112)).shape) |
|
|
