File size: 1,916 Bytes
902df3d | 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 | import torch, torch.nn as nn
class DW(nn.Module): # depthwise-separable bottleneck (MobileFaceNet style)
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) # training-only: map to teacher space
s.cls=nn.Linear(emb,n_cls) # training-only
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)
|