Skip to content

Instantly share code, notes, and snippets.

@shawwn

shawwn/ConvMixr.py

Created October 7, 2021 10:17
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save shawwn/5ce14c8f95953c02ee997d39cd31eed2 to your computer and use it in GitHub Desktop.
Save shawwn/5ce14c8f95953c02ee997d39cd31eed2 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
ConvMixr.py
def ConvMixr(h,d,k,p,n):
def A(x):
return Sequential(x, GELU(), BatchNorm2d(h))
class R(Sequential):
def forward(self, x):
return self[0](x) + x
return Sequential(
A(Conv2d(3,h,p,p)),
*[
Sequential(
R(A(Conv2d(h,h,k,groups=h,padding=k//2))),
A(Conv2d(h,h,1)))
for i in range(d)],
AdaptiveAvgPool2d((1,1)),
Flatten(),
Linear(h,n))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment