ConvMixer

代码：

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.GEL