paper
import torch.nn as nn
import torch
class sMLPBlock(nn.Module):
'''
稀疏MLP 不是一个样本的所有特征通过全连接层 而是部分通过全连接层
'''
def __init__(self, W, H, channels):
super().__init__()
assert W == H
self.channels = channels
self.activation = nn.GELU()
self.BN = nn.BatchNorm2d(channels)
self.proj_h = nn.Conv2d(H, H, (1, 1))
self.proh_w = nn.Conv2d(W, W, (1, 1))
self.fuse = nn.Conv2d(channels*3, channels, (1,1), (1,1), bias=False)
# 也可以这样写
# self.proj_h2=nn.Linear(H,H)
# self.proh_w2=nn.Linear(W,W)
# self.fuse2=nn.Linear(channels*3,channels)
def forward(self, x):
x = self.activation(self.BN(x))
x_w = self.proj_h(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
x_h = self.proh_w(x.permute(0, 2, 1, 3)).permute(0, 2, 1, 3)
x = self.fuse(torch.cat([x, x_h, x_w], dim=1))
print(x)
# 写法2
# x_h2=self.proj_h2(x.permute(0,1,3,2)).permute(0,1,3,2)
# x_w2=self.proh_w2(x)
# x_fuse=torch.cat([x_h2,x_w2,x],dim=1)
# x_2=self.fuse2(x_fuse.permute(0,2,3,1)).permute(0,3,1,2)
# print(x_2)
return x
if __name__ == '__main__':
x = torch.randn(1, 3, 2, 2).cuda() # 输入 B C H W
model = sMLPBlock(2,2,3).cuda()
res=model(x)
