当前输入维度为[B,C,H,W]，那么应该如何转换为import torch
import torch.nn as nn
from spikingjelly.clock_driven.neuron import MultiStepLIFNode
from spikingjelly.clock_driven import layer, surrogate
from timm.models.layers import to_2tuple, trunc_normal_, DropPath
from timm.models.registry import register_model
from timm.models.vision_transformer import cfg
from einops.layers.torch import Rearrange
import torch.nn.functional as F
from functools import partial

all = [‘spikformer’]

class MLP(nn.Module):
def init(self, in_features, hidden_features=None, out_features=None, drop=0.):
super().init()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1_conv = nn.Conv1d(in_features, hidden_features, kernel_size=1, stride=1)
self.fc1_bn = nn.BatchNorm1d(hidden_features)
self.fc1_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)

self.fc2_conv = nn.Conv1d(hidden_features, out_features, kernel_size=1, stride=1)
self.fc2_bn = nn.BatchNorm1d(out_features)
self.fc2_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)
self.c_hidden = hidden_features
self.c_output = out_features
def forward(self, x):
T,B,C,N = x.shape
x = self.fc1_conv(x.flatten(0,1))
x = self.fc1_bn(x).reshape(T,B,self.c_hidden,N).contiguous()
x = self.fc1_lif(x)

x = self.fc2_conv(x.flatten(0,1))
x = self.fc2_bn(x).reshape(T,B,C,N).contiguous()
x = self.fc2_lif(x)
return x

class SSA(nn.Module):
def init(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
super().init()
assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}.“
self.dim = dim
self.num_heads = num_heads
self.scale = 0.25

self.q_conv = nn.Conv1d(dim, dim, kernel_size=1, stride=1,bias=False)

self.q_bn = nn.BatchNorm1d(dim)
self.q_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)

self.k_conv = nn.Conv1d(dim, dim, kernel_size=1, stride=1,bias=False)
self.k_bn = nn.BatchNorm1d(dim)
self.k_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)

self.v_conv = nn.Conv1d(dim, dim, kernel_size=1, stride=1,bias=False)
self.v_bn = nn.BatchNorm1d(dim)
self.v_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)

self.attn_drop = nn.Dropout(0.2)
self.res_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)
self.attn_lif = MultiStepLIFNode(tau=2.0, v_threshold=0.5, detach_reset=True, backend=‘cupy’)

self.proj_conv = nn.Conv1d(dim, dim, kernel_size=1, stride=1)
self.proj_bn = nn.BatchNorm1d(dim)
self.proj_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)

def forward(self, x):
T, B, C, N = x.shape
x_for_qkv = x.flatten(0, 1)
q_conv_out = self.q_conv(x_for_qkv)
q_conv_out = self.q_bn(q_conv_out).reshape(T,B,C,N).contiguous()
q_conv_out = self.q_lif(q_conv_out)
q = q_conv_out.transpose(-1, -2).reshape(T, B, N, self.num_heads, C//self.num_heads).permute(0, 1, 3, 2, 4).contiguous()

k_conv_out = self.k_conv(x_for_qkv)
k_conv_out = self.k_bn(k_conv_out).reshape(T,B,C,N).contiguous()
k_conv_out = self.k_lif(k_conv_out)
k = k_conv_out.transpose(-1, -2).reshape(T, B, N, self.num_heads, C//self.num_heads).permute(0, 1, 3, 2, 4).contiguous()

v_conv_out = self.v_conv(x_for_qkv)
v_conv_out = self.v_bn(v_conv_out).reshape(T,B,C,N).contiguous()
v_conv_out = self.v_lif(v_conv_out)
v = v_conv_out.transpose(-1, -2).reshape(T, B, N, self.num_heads, C//self.num_heads).permute(0, 1, 3, 2, 4).contiguous()

attn = (q @ k.transpose(-2, -1))
x = (attn @ v) * self.scale

x = x.transpose(3, 4).reshape(T, B, C, N).contiguous()
x = self.attn_lif(x)
x = x.flatten(0,1)
x = self.proj_lif(self.proj_bn(self.proj_conv(x)).reshape(T,B,C,N))

return x


class Block(nn.Module):
def init(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., norm_layer=nn.LayerNorm, sr_ratio=1):
super().init()
self.norm1 = norm_layer(dim)
self.attn = SSA(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop)

def forward(self, x):
x = x + self.attn(x)
x = x + (self.mlp(x))
return x


class SPS(nn.Module):
def init(self, img_size_h=128, img_size_w=128, patch_size=4, in_channels=2, embed_dims=256):
super().init()
self.image_size = [img_size_h, img_size_w]
patch_size = to_2tuple(patch_size)
self.patch_size = patch_size
self.C = in_channels
self.H, self.W = self.image_size[0] // patch_size[0], self.image_size[1] // patch_size[1]
self.num_patches = self.H * self.W
self.proj_conv = nn.Conv2d(in_channels, embed_dims//8, kernel_size=3, stride=1, padding=1, bias=False)
self.proj_bn = nn.BatchNorm2d(embed_dims//8)
self.proj_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)
self.maxpool = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

self.proj_conv1 = nn.Conv2d(embed_dims//8, embed_dims//4, kernel_size=3, stride=1, padding=1, bias=False)
self.proj_bn1 = nn.BatchNorm2d(embed_dims//4)
self.proj_lif1 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)
self.maxpool1 = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

self.proj_conv2 = nn.Conv2d(embed_dims//4, embed_dims//2, kernel_size=3, stride=1, padding=1, bias=False)
self.proj_bn2 = nn.BatchNorm2d(embed_dims//2)
self.proj_lif2 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)
self.maxpool2 = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

self.proj_conv3 = nn.Conv2d(embed_dims//2, embed_dims, kernel_size=3, stride=1, padding=1, bias=False)
self.proj_bn3 = nn.BatchNorm2d(embed_dims)
self.proj_lif3 = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)
self.maxpool3 = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

self.rpe_conv = nn.Conv2d(embed_dims, embed_dims, kernel_size=3, stride=1, padding=1, bias=False)
self.rpe_bn = nn.BatchNorm2d(embed_dims)
self.rpe_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend=‘cupy’)
def forward(self, x):
T, B, C, H, W = x.shape
x = self.proj_conv(x.flatten(0, 1)) # have some fire value
x = self.proj_bn(x).reshape(T,B,-1,H,W).contiguous()
x = self.proj_lif(x).flatten(0,1).contiguous()
x = self.maxpool(x)

x = self.proj_conv1(x)
x = self.proj_bn1(x).reshape(T, B, -1, 64, 64).contiguous()
x = self.proj_lif1(x).flatten(0, 1).contiguous()
x = self.maxpool1(x)

x = self.proj_conv2(x)
x = self.proj_bn2(x).reshape(T, B, -1, 32, 32).contiguous()
x = self.proj_lif2(x).flatten(0, 1).contiguous()
x = self.maxpool2(x)

x = self.proj_conv3(x)
x = self.proj_bn3(x).reshape(T, B, -1, 16, 16).contiguous()
x = self.proj_lif3(x).flatten(0, 1).contiguous()
x = self.maxpool3(x)

x_rpe = self.rpe_bn(self.rpe_conv(x)).reshape(T,B,-1,H//16,W//16).contiguous()
x_rpe = self.rpe_lif(x_rpe).flatten(0,1)
x = x + x_rpe
x = x.reshape(T, B, -1, (H//16)*(H//16)).contiguous()

return x

class Spikformer(nn.Module):
def init(self,
img_size_h=128, img_size_w=128, patch_size=16, in_channels=2, num_classes=11,
embed_dims=[64, 128, 256], num_heads=[1, 2, 4], mlp_ratios=[4, 4, 4], qkv_bias=False, qk_scale=None,
drop_rate=0., attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
depths=[6, 8, 6], sr_ratios=[8, 4, 2]
):
super().init()
self.num_classes = num_classes
self.depths = depths

dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depths)] # stochastic depth decay rule

patch_embed = SPS(img_size_h=img_size_h,
img_size_w=img_size_w,
patch_size=patch_size,
in_channels=in_channels,
embed_dims=embed_dims)
num_patches = patch_embed.num_patches
pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dims))
block = nn.ModuleList([Block(
dim=embed_dims, num_heads=num_heads, mlp_ratio=mlp_ratios, qkv_bias=qkv_bias,
qk_scale=qk_scale, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[j],
norm_layer=norm_layer, sr_ratio=sr_ratios)
for j in range(depths)])

setattr(self, f"patch_embed”, patch_embed)
setattr(self, f"pos_embed", pos_embed)
setattr(self, f"block", block)

self.norm = norm_layer(embed_dims)
# classification head 这里不需要脉冲，因为输入的是在T时长平均发射值
self.head = nn.Linear(embed_dims, num_classes) if num_classes > 0 else nn.Identity()

pos_embed = getattr(self, f"pos_embed")
trunc_normal(pos_embed, std=.02)
self.apply(self.init_weights)

@torch.jit.ignore
def get_pos_embed(self, pos_embed, patch_embed, H, W):
if H * W == self.patch_embed1.num_patches:
return pos_embed
else:
return F.interpolate(
pos_embed.reshape(1, patch_embed.H, patch_embed.W, -1).permute(0, 3, 1, 2),
size=(H, W), mode=“bilinear”).reshape(1, -1, H * W).permute(0, 2, 1)

def init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant(m.bias, 0)
nn.init.constant_(m.weight, 1.0)

def forward_features(self, x):

block = getattr(self, f"block")
patch_embed = getattr(self, f"patch_embed")
x = patch_embed(x)
for blk in block:
x = blk(x)
return x.mean(3)

def forward(self, x):
x = x.permute(1, 0, 2, 3, 4) # [T, N, 2, *, *]
x = self.forward_features(x)
x = self.head(x.mean(0))
return x


@register_model
def spikformer(pretrained=False, **kwargs):
model = Spikformer(
patch_size=16, embed_dims=256, num_heads=16, mlp_ratios=4,
in_channels=2, num_classes=10, qkv_bias=False,
norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=2, sr_ratios=1,
**kwargs
)
model.default_cfg = _cfg()
return model



class Multi_Context(nn.Module):
def init(self, input_channels, output_channels): # 修正了def init为def init
super(Multi_Context, self).init()
self.linear1 = nn.Linear(input_channels, output_channels)
self.linear2 = nn.Linear(input_channels, output_channels)
self.linear3 = nn.Linear(input_channels, output_channels)
self.linear_final = nn.Linear(output_channels * 3, output_channels)

def forward(self, x):
x1 = F.relu(self.linear1(x))
x2 = F.relu(self.linear2(x))
x3 = F.relu(self.linear3(x))
x = torch.cat([x1, x2, x3], dim=-1) # Concatenate along the feature dimension
x = self.linear_final(x)
return x模型期望的输入？