代码:
import torch.nn as nn
import torch
class Residual(nn.Module):
def __init__(self,fn):
super().__init__()
self.fn=fn
def forward(self,x):
return x+self.fn(x)
def ConvMixer(dim,depth,kernel_size=5,patch_size=7,num_classes=1000):
return nn.Sequential(
nn.Conv2d(3,dim,kernel_size=patch_size,stride=patch_size),
nn.GELU(),
nn.BatchNorm2d(dim),
*[nn.Sequential(
Residual(nn.Sequential(
nn.Conv2d(dim,dim,kernel_size=kernel_size,groups=dim,padding=kernel_size//2),
nn.GELU(),
nn.BatchNorm2d(dim)
)),
nn.Conv2d(dim,dim,kernel_size=1),
nn.GELU(),
nn.BatchNorm2d(dim)
) for _ in range(depth)],
nn.AdaptiveAvgPool2d(1),
nn.