使用SS2D写了一个简单的神经网络进行训练，但是训练报错：
NotImplementedError: You must implement either the backward or vjp method for your custom autograd.Function to use it with backward mode AD.
环境：
CUDA11.8
torch=2.0.0
mamba_ssm=2.0.2
causal-conv1d=1.2.1
Ubuntu 22.04 , Python3.11.8
python
import torch
import torch.nn as nn
from functools import partial
from typing import Callable
import math
from einops import repeat
from timm.models.layers import DropPath
from mmyolo.registry import MODELS # new
try:
from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, selective_scan_ref
except Exception as e:
# print(e, flush=True)
pass
try:
"sscore acts the same as mamba_ssm"
SSMODE = "sscore"
import selective_scan_cuda_core
except Exception as e:
print(e, flush=True)
"you should install mamba_ssm to use this"
SSMODE = "mamba_ssm"
import selective_scan_cuda
class SelectiveScanMamba(torch.autograd.Function):
# comment all checks if inside cross_selective_scan
@staticmethod
@torch.cuda.amp.custom_fwd
def forward(ctx, u, delta, A, B, C, D=None, delta_bias=None, delta_softplus=False, nrows=1, backnrows=1,
oflex=True):
ctx.delta_softplus = delta_softplus
out, x, *rest = selective_scan_cuda.fwd(u, delta, A, B, C, D, None, delta_bias, delta_softplus)
ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)
return out
@staticmethod
@torch.cuda.amp.custom_bwd
def backward(ctx, dout, *args):
u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensors
if dout.stride(-1) != 1:
dout = dout.contiguous()
du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda.bwd(
u, delta, A, B, C, D, None, delta_bias, dout, x, None, None, ctx.delta_softplus,
False
)
return (du, ddelta, dA, dB, dC, dD, ddelta_bias, None, None, None, None)
class SelectiveScanCore(torch.autograd.Function):
@staticmethod
@torch.cuda.amp.custom_fwd
def forward(ctx, u, delta, A, B, C, D=None, delta_bias=None, delta_softplus=False, nrows=1, backnrows=1,
oflex=True):
ctx.delta_softplus = delta_softplus
if SSMODE == "mamba_ssm":
out, x, *rest = selective_scan_cuda.fwd(u, delta, A, B, C, D, None, delta_bias, delta_softplus)
else:
out, x, *rest = selective_scan_cuda_core.fwd(u, delta, A, B, C, D, delta_bias, delta_softplus, nrows)
ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)
# out, x, *rest = selective_scan_cuda_core.fwd(u, delta, A, B, C, D, delta_bias, delta_softplus, 1)
# ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)
return out
@staticmethod
@torch.cuda.amp.custom_bwd
def backward(ctx, dout, *args):
u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensors
if dout.stride(-1) != 1:
dout = dout.contiguous()
if SSMODE == "mamba_ssm":
du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda.bwd(
u, delta, A, B, C, D, None, delta_bias, dout, x, None, None, ctx.delta_softplus,
False # option to recompute out_z, not used here
)
else:
du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda_core.bwd(
u, delta, A, B, C, D, delta_bias, dout, x, ctx.delta_softplus, 1
)
return (du, ddelta, dA, dB, dC, dD, ddelta_bias, None, None, None, None)
class SelectiveScanFake(torch.autograd.Function):
# comment all checks if inside cross_selective_scan
@staticmethod
@torch.cuda.amp.custom_fwd
def forward(ctx, u, delta, A, B, C, D=None, delta_bias=None, delta_softplus=False, nrows=1, backnrows=1,
oflex=True):
ctx.delta_softplus = delta_softplus
ctx.backnrows = backnrows
x = delta
out = u
ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)
return out
@staticmethod
@torch.cuda.amp.custom_bwd
def backward(ctx, dout, *args):
u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensors
if dout.stride(-1) != 1:
dout = dout.contiguous()
du, ddelta, dA, dB, dC, dD, ddelta_bias = u * 0, delta * 0, A * 0, B * 0, C * 0, C * 0, (
D * 0 if D else None), (delta_bias * 0 if delta_bias else None)
return (du, ddelta, dA, dB, dC, dD, ddelta_bias, None, None, None, None)
class CrossScan(torch.autograd.Function):
@staticmethod
def forward(ctx, x: torch.Tensor):
B, C, H, W = x.shape
ctx.shape = (B, C, H, W)
# xs = x.new_empty((B, 4, C, H * W))
xs = x.new_empty((B, 8, C, H * W))
xs[:, 0] = x.flatten(2, 3)
xs[:, 1] = x.transpose(dim0=2, dim1=3).flatten(2, 3)
xs[:, 2:4] = torch.flip(xs[:, 0:2], dims=[-1])
xs[:, 4] = diagonal_gather(x)
xs[:, 5] = antidiagonal_gather(x)
xs[:, 6:8] = torch.flip(xs[:, 4:6], dims=[-1])
return xs
@staticmethod
def backward(ctx, ys: torch.Tensor):
# out: (b, k, d, l)
B, C, H, W = ctx.shape
L = H * W
# ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, -1, L)
y_rb = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, -1, L)
# y = ys[:, 0] + ys[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, -1, L)
y_rb = y_rb[:, 0] + y_rb[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, -1, L)
y_rb = y_rb.view(B, -1, H, W)
y_da = ys[:, 4:6] + ys[:, 6:8].flip(dims=[-1]).view(B, 2, -1, L)
y_da = diagonal_scatter(y_da[:, 0], (B, C, H, W)) + antidiagonal_scatter(y_da[:, 1], (B, C, H, W))
y_res = y_rb + y_da
# return y.view(B, -1, H, W)
return y_res
def antidiagonal_gather(tensor):
B, C, H, W = tensor.size()
shift = torch.arange(H, device=tensor.device).unsqueeze(1)
index = (torch.arange(W, device=tensor.device) - shift) % W
expanded_index = index.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
return tensor.gather(3, expanded_index).transpose(-1, -2).reshape(B, C, H * W)
def diagonal_gather(tensor):
B, C, H, W = tensor.size()
shift = torch.arange(H, device=tensor.device).unsqueeze(1)
index = (shift + torch.arange(W, device=tensor.device)) % W
expanded_index = index.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
return tensor.gather(3, expanded_index).transpose(-1, -2).reshape(B, C, H * W)
def diagonal_scatter(tensor_flat, original_shape):
B, C, H, W = original_shape
shift = torch.arange(H, device=tensor_flat.device).unsqueeze(1)
index = (shift + torch.arange(W, device=tensor_flat.device)) % W
expanded_index = index.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
result_tensor = torch.zeros(B, C, H, W, device=tensor_flat.device, dtype=tensor_flat.dtype)
tensor_reshaped = tensor_flat.reshape(B, C, W, H).transpose(-1, -2)
result_tensor.scatter_(3, expanded_index, tensor_reshaped)
return result_tensor
def antidiagonal_scatter(tensor_flat, original_shape):
B, C, H, W = original_shape
shift = torch.arange(H, device=tensor_flat.device).unsqueeze(1)
index = (torch.arange(W, device=tensor_flat.device) - shift) % W
expanded_index = index.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
result_tensor = torch.zeros(B, C, H, W, device=tensor_flat.device, dtype=tensor_flat.dtype)
tensor_reshaped = tensor_flat.reshape(B, C, W, H).transpose(-1, -2)
result_tensor.scatter_(3, expanded_index, tensor_reshaped)
return result_tensor
class CrossMerge(torch.autograd.Function):
@staticmethod
def forward(ctx, ys: torch.Tensor):
B, K, D, H, W = ys.shape
ctx.shape = (H, W)
ys = ys.view(B, K, D, -1)
# ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
# y = ys[:, 0] + ys[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, D, -1)
y_rb = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
y_rb = y_rb[:, 0] + y_rb[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, D, -1)
y_rb = y_rb.view(B, -1, H, W)
y_da = ys[:, 4:6] + ys[:, 6:8].flip(dims=[-1]).view(B, 2, D, -1)
y_da = diagonal_scatter(y_da[:, 0], (B, D, H, W)) + antidiagonal_scatter(y_da[:, 1], (B, D, H, W))
y_res = y_rb + y_da
return y_res.view(B, D, -1)
# return y
class CrossScan_Ab_1direction(torch.autograd.Function):
@staticmethod
def forward(ctx, x: torch.Tensor):
B, C, H, W = x.shape
ctx.shape = (B, C, H, W)
xs = x.view(B, 1, C, H * W).repeat(1, 4, 1, 1).contiguous()
return xs
@staticmethod
def backward(ctx, ys: torch.Tensor):
# out: (b, k, d, l)
B, C, H, W = ctx.shape
y = ys.sum(dim=1).view(B, C, H, W)
return y
class CrossMerge_Ab_1direction(torch.autograd.Function):
@staticmethod
def forward(ctx, ys: torch.Tensor):
B, K, D, H, W = ys.shape
ctx.shape = (H, W)
y = ys.sum(dim=1).view(B, D, H * W)
return y
@staticmethod
def backward(ctx, x: torch.Tensor):
# B, D, L = x.shape
# out: (b, k, d, l)
H, W = ctx.shape
B, C, L = x.shape
xs = x.view(B, 1, C, L).repeat(1, 4, 1, 1).contiguous().view(B, 4, C, H, W)
return xs
class CrossScan_Ab_2direction(torch.autograd.Function):
@staticmethod
def forward(ctx, x: torch.Tensor):
B, C, H, W = x.shape
ctx.shape = (B, C, H, W)
xs = x.new_empty((B, 4, C, H * W))
xs[:, 0] = x.flatten(2, 3)
xs[:, 1] = x.flatten(2, 3)
xs[:, 2:4] = torch.flip(xs[:, 0:2], dims=[-1])
return xs
@staticmethod
def backward(ctx, ys: torch.Tensor):
# out: (b, k, d, l)
B, C, H, W = ctx.shape
L = H * W
ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, -1, L)
y = ys[:, 0] + ys[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, -1, L)
return y.view(B, -1, H, W)
class CrossMerge_Ab_2direction(torch.autograd.Function):
@staticmethod
def forward(ctx, ys: torch.Tensor):
B, K, D, H, W = ys.shape
ctx.shape = (H, W)
ys = ys.view(B, K, D, -1)
ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
y = ys.sum(dim=1)
return y
@staticmethod
def backward(ctx, x: torch.Tensor):
# B, D, L = x.shape
# out: (b, k, d, l)
H, W = ctx.shape
B, C, L = x.shape
xs = x.new_empty((B, 4, C, L))
xs[:, 0] = x
xs[:, 1] = x
xs[:, 2:4] = torch.flip(xs[:, 0:2], dims=[-1])
xs = xs.view(B, 4, C, H, W)
return xs
def cross_selective_scan(
x: torch.Tensor = None,
x_proj_weight: torch.Tensor = None,
x_proj_bias: torch.Tensor = None,
dt_projs_weight: torch.Tensor = None,
dt_projs_bias: torch.Tensor = None,
A_logs: torch.Tensor = None,
Ds: torch.Tensor = None,
delta_softplus=True,
out_norm: torch.nn.Module = None,
out_norm_shape="v0",
# ==============================
to_dtype=True, # True: final out to dtype
force_fp32=False, # True: input fp32
# ==============================
nrows=-1, # for SelectiveScanNRow; 0: auto; -1: disable;
backnrows=-1, # for SelectiveScanNRow; 0: auto; -1: disable;
ssoflex=True, # True: out fp32 in SSOflex; else, SSOflex is the same as SSCore
# ==============================
SelectiveScan=None,
CrossScan=CrossScan,
CrossMerge=CrossMerge,
):
# out_norm: whatever fits (B, L, C); LayerNorm; Sigmoid; Softmax(dim=1);...
B, D, H, W = x.shape
D, N = A_logs.shape
K, D, R = dt_projs_weight.shape
L = H * W
if nrows == 0:
if D % 4 == 0:
nrows = 4
elif D % 3 == 0:
nrows = 3
elif D % 2 == 0:
nrows = 2
else:
nrows = 1
if backnrows == 0:
if D % 4 == 0:
backnrows = 4
elif D % 3 == 0:
backnrows = 3
elif D % 2 == 0:
backnrows = 2
else:
backnrows = 1
def selective_scan(u, delta, A, B, C, D=None, delta_bias=None, delta_softplus=True):
return SelectiveScan.apply(u, delta, A, B, C, D, delta_bias, delta_softplus, nrows, backnrows, ssoflex)
xs = CrossScan.apply(x)
x_dbl = torch.einsum("b k d l, k c d -> b k c l", xs, x_proj_weight)
if x_proj_bias is not None:
x_dbl = x_dbl + x_proj_bias.view(1, K, -1, 1)
dts, Bs, Cs = torch.split(x_dbl, [R, N, N], dim=2)
dts = torch.einsum("b k r l, k d r -> b k d l", dts, dt_projs_weight)
xs = xs.view(B, -1, L)
dts = dts.contiguous().view(B, -1, L)
As = -torch.exp(A_logs.to(torch.float)) # (k * c, d_state)
Bs = Bs.contiguous()
Cs = Cs.contiguous()
Ds = Ds.to(torch.float) # (K * c)
delta_bias = dt_projs_bias.view(-1).to(torch.float)
if force_fp32:
xs = xs.to(torch.float)
dts = dts.to(torch.float)
Bs = Bs.to(torch.float)
Cs = Cs.to(torch.float)
ys: torch.Tensor = selective_scan(
xs, dts, As, Bs, Cs, Ds, delta_bias, delta_softplus
).view(B, K, -1, H, W)
y: torch.Tensor = CrossMerge.apply(ys)
if out_norm_shape in ["v1"]: # (B, C, H, W)
y = out_norm(y.view(B, -1, H, W)).permute(0, 2, 3, 1) # (B, H, W, C)
else: # (B, L, C)
y = y.transpose(dim0=1, dim1=2).contiguous() # (B, L, C)
y = out_norm(y).view(B, H, W, -1)
return (y.to(x.dtype) if to_dtype else y)
class CrossMerge_Ab_1direction(torch.autograd.Function):
@staticmethod
def forward(ctx, ys: torch.Tensor):
B, K, D, H, W = ys.shape
ctx.shape = (H, W)
y = ys.sum(dim=1).view(B, D, H * W)
return y
@staticmethod
def backward(ctx, x: torch.Tensor):
# B, D, L = x.shape
# out: (b, k, d, l)
H, W = ctx.shape
B, C, L = x.shape
xs = x.view(B, 1, C, L).repeat(1, 4, 1, 1).contiguous().view(B, 4, C, H, W)
return xs
class SS2D(nn.Module):
def __init__(
self,
# basic dims ===========
d_model=64,
d_state=16,
ssm_ratio=2.0,
dt_rank="auto",
act_layer=nn.SiLU,
# dwconv ===============
d_conv=3, # < 2 means no conv
conv_bias=True,
# ======================
dropout=0.0,
bias=False,
# dt init ==============
dt_min=0.001,
dt_max=0.1,
dt_init="random",
dt_scale=1.0,
dt_init_floor=1e-4,
initialize="v0",
# ======================
forward_type="v2",
# ======================
**kwargs,
):
factory_kwargs = {"device": None, "dtype": None}
super().__init__()
d_inner = int(ssm_ratio * d_model)
dt_rank = math.ceil(d_model / 16) if dt_rank == "auto" else dt_rank
self.d_conv = d_conv
# tags for forward_type ==============================
def checkpostfix(tag, value):
ret = value[-len(tag):] == tag
if ret:
value = value[:-len(tag)]
return ret, value
self.disable_force32, forward_type = checkpostfix("no32", forward_type)
self.disable_z, forward_type = checkpostfix("noz", forward_type)
self.disable_z_act, forward_type = checkpostfix("nozact", forward_type)
# softmax | sigmoid | dwconv | norm ===========================
if forward_type[-len("none"):] == "none":
forward_type = forward_type[:-len("none")]
self.out_norm = nn.Identity()
elif forward_type[-len("dwconv3"):] == "dwconv3":
forward_type = forward_type[:-len("dwconv3")]
self.out_norm = nn.Conv2d(d_inner, d_inner, kernel_size=3, padding=1, groups=d_inner, bias=False)
self.out_norm_shape = "v1"
elif forward_type[-len("softmax"):] == "softmax":
forward_type = forward_type[:-len("softmax")]
self.out_norm = nn.Softmax(dim=1)
elif forward_type[-len("sigmoid"):] == "sigmoid":
forward_type = forward_type[:-len("sigmoid")]
self.out_norm = nn.Sigmoid()
else:
self.out_norm = nn.LayerNorm(d_inner)
# forward_type debug =======================================
FORWARD_TYPES = dict(
v0=self.forward_corev0,
v2=partial(self.forward_corev2, force_fp32=True, SelectiveScan=SelectiveScanCore),
fake=partial(self.forward_corev2, force_fp32=(not self.disable_force32), SelectiveScan=SelectiveScanFake),
# v1=partial(self.forward_corev2, force_fp32=True, SelectiveScan=SelectiveScanOflex),
v01=partial(self.forward_corev2, force_fp32=(not self.disable_force32), SelectiveScan=SelectiveScanMamba),
)
self.forward_core = FORWARD_TYPES.get(forward_type, None)
k_group = 8 if forward_type not in ["debugscan_sharessm"] else 1
# in proj =======================================
d_proj = d_inner if self.disable_z else (d_inner * 2)
self.in_proj = nn.Linear(d_model, d_proj, bias=bias, **factory_kwargs)
self.act: nn.Module = act_layer()
# conv =======================================
if d_conv > 1:
self.conv2d = nn.Conv2d(
in_channels=d_inner,
out_channels=d_inner,
groups=d_inner,
bias=conv_bias,
kernel_size=d_conv,
padding=(d_conv - 1) // 2,
**factory_kwargs,
)
# x proj ============================
self.x_proj = [
nn.Linear(d_inner, (dt_rank + d_state * 2), bias=False, **factory_kwargs)
for _ in range(k_group)
]
self.x_proj_weight = nn.Parameter(torch.stack([t.weight for t in self.x_proj], dim=0)) # (K, N, inner)
del self.x_proj
# out proj =======================================
self.out_proj = nn.Linear(d_inner, d_model, bias=bias, **factory_kwargs)
self.dropout = nn.Dropout(dropout) if dropout > 0. else nn.Identity()
if initialize in ["v0"]:
# dt proj ============================
self.dt_projs = [
self.dt_init(dt_rank, d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor, **factory_kwargs)
for _ in range(k_group)
]
self.dt_projs_weight = nn.Parameter(
torch.stack([t.weight for t in self.dt_projs], dim=0)) # (K, inner, rank)
self.dt_projs_bias = nn.Parameter(torch.stack([t.bias for t in self.dt_projs], dim=0)) # (K, inner)
del self.dt_projs
# A, D =======================================
self.A_logs = self.A_log_init(d_state, d_inner, copies=k_group, merge=True) # (K * D, N)
self.Ds = self.D_init(d_inner, copies=k_group, merge=True) # (K * D)
elif initialize in ["v1"]:
# simple init dt_projs, A_logs, Ds
self.Ds = nn.Parameter(torch.ones((k_group * d_inner)))
self.A_logs = nn.Parameter(
torch.randn((k_group * d_inner, d_state))) # A == -A_logs.exp() < 0; # 0 < exp(A * dt) < 1
self.dt_projs_weight = nn.Parameter(torch.randn((k_group, d_inner, dt_rank)))
self.dt_projs_bias = nn.Parameter(torch.randn((k_group, d_inner)))
elif initialize in ["v2"]:
# simple init dt_projs, A_logs, Ds
self.Ds = nn.Parameter(torch.ones((k_group * d_inner)))
self.A_logs = nn.Parameter(
torch.zeros((k_group * d_inner, d_state))) # A == -A_logs.exp() < 0; # 0 < exp(A * dt) < 1
self.dt_projs_weight = nn.Parameter(torch.randn((k_group, d_inner, dt_rank)))
self.dt_projs_bias = nn.Parameter(torch.randn((k_group, d_inner)))
@staticmethod
def dt_init(dt_rank, d_inner, dt_scale=1.0, dt_init="random", dt_min=0.001, dt_max=0.1, dt_init_floor=1e-4,
**factory_kwargs):
dt_proj = nn.Linear(dt_rank, d_inner, bias=True, **factory_kwargs)
# Initialize special dt projection to preserve variance at initialization
dt_init_std = dt_rank ** -0.5 * dt_scale
if dt_init == "constant":
nn.init.constant_(dt_proj.weight, dt_init_std)
elif dt_init == "random":
nn.init.uniform_(dt_proj.weight, -dt_init_std, dt_init_std)
else:
raise NotImplementedError
# Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max
dt = torch.exp(
torch.rand(d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
+ math.log(dt_min)
).clamp(min=dt_init_floor)
# Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
inv_dt = dt + torch.log(-torch.expm1(-dt))
with torch.no_grad():
dt_proj.bias.copy_(inv_dt)
# Our initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
# dt_proj.bias._no_reinit = True
return dt_proj
@staticmethod
def A_log_init(d_state, d_inner, copies=-1, device=None, merge=True):
# S4D real initialization
A = repeat(
torch.arange(1, d_state + 1, dtype=torch.float32, device=device),
"n -> d n",
d=d_inner,
).contiguous()
A_log = torch.log(A) # Keep A_log in fp32
if copies > 0:
A_log = repeat(A_log, "d n -> r d n", r=copies)
if merge:
A_log = A_log.flatten(0, 1)
A_log = nn.Parameter(A_log)
A_log._no_weight_decay = True
return A_log
@staticmethod
def D_init(d_inner, copies=-1, device=None, merge=True):
# D "skip" parameter
D = torch.ones(d_inner, device=device)
if copies > 0:
D = repeat(D, "n1 -> r n1", r=copies)
if merge:
D = D.flatten(0, 1)
D = nn.Parameter(D) # Keep in fp32
D._no_weight_decay = True
return D
# only used to run previous version
def forward_corev0(self, x: torch.Tensor, to_dtype=False, channel_first=False):
def selective_scan(u, delta, A, B, C, D=None, delta_bias=None, delta_softplus=True, nrows=1):
return SelectiveScanCore.apply(u, delta, A, B, C, D, delta_bias, delta_softplus, nrows, False)
if not channel_first:
x = x.permute(0, 3, 1, 2).contiguous()
B, D, H, W = x.shape
D, N = self.A_logs.shape
K, D, R = self.dt_projs_weight.shape
L = H * W
x_hwwh = torch.stack([x.view(B, -1, L), torch.transpose(x, dim0=2, dim1=3).contiguous().view(B, -1, L)],
dim=1).view(B, 2, -1, L)
xs = torch.cat([x_hwwh, torch.flip(x_hwwh, dims=[-1])], dim=1) # (b, k, d, l)
x_dbl = torch.einsum("b k d l, k c d -> b k c l", xs, self.x_proj_weight)
# x_dbl = x_dbl + self.x_proj_bias.view(1, K, -1, 1)
dts, Bs, Cs = torch.split(x_dbl, [R, N, N], dim=2)
dts = torch.einsum("b k r l, k d r -> b k d l", dts, self.dt_projs_weight)
xs = xs.float().view(B, -1, L) # (b, k * d, l)
dts = dts.contiguous().float().view(B, -1, L) # (b, k * d, l)
Bs = Bs.float() # (b, k, d_state, l)
Cs = Cs.float() # (b, k, d_state, l)
As = -torch.exp(self.A_logs.float()) # (k * d, d_state)
Ds = self.Ds.float() # (k * d)
dt_projs_bias = self.dt_projs_bias.float().view(-1) # (k * d)
out_y = selective_scan(
xs, dts,
As, Bs, Cs, Ds,
delta_bias=dt_projs_bias,
delta_softplus=True,
).view(B, K, -1, L)
# assert out_y.dtype == torch.float
inv_y = torch.flip(out_y[:, 2:4], dims=[-1]).view(B, 2, -1, L)
wh_y = torch.transpose(out_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
invwh_y = torch.transpose(inv_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
y = out_y[:, 0] + inv_y[:, 0] + wh_y + invwh_y
y = y.transpose(dim0=1, dim1=2).contiguous() # (B, L, C)
y = self.out_norm(y).view(B, H, W, -1)
return (y.to(x.dtype) if to_dtype else y)
def forward_corev2(self, x: torch.Tensor, channel_first=False, SelectiveScan=SelectiveScanCore,
cross_selective_scan=cross_selective_scan, force_fp32=None):
if not channel_first:
x = x.permute(0, 3, 1, 2).contiguous()
x = cross_selective_scan(
x, self.x_proj_weight, None, self.dt_projs_weight, self.dt_projs_bias,
self.A_logs, self.Ds, delta_softplus=True,
out_norm=getattr(self, "out_norm", None),
out_norm_shape=getattr(self, "out_norm_shape", "v0"),
force_fp32=force_fp32,
SelectiveScan=SelectiveScan,
)
return x
def forward(self, x: torch.Tensor, **kwargs):
with_dconv = (self.d_conv > 1) # True
x = self.in_proj(x)
if not self.disable_z:
x, z = x.chunk(2, dim=-1) # (b, h, w, d)
if not self.disable_z_act:
z = self.act(z)
if with_dconv:
x = x.permute(0, 3, 1, 2).contiguous()
x = self.conv2d(x) # (b, d, h, w)
x = self.act(x)
y = self.forward_core(x, channel_first=with_dconv)
if not self.disable_z:
y = y * z # 元素级乘法
out = self.dropout(self.out_proj(y))
return out
class AttentionBlockWithMLP(nn.Module):
def __init__(self, in_channels, reduction=16):
super(AttentionBlockWithMLP, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(in_channels, in_channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(in_channels // reduction, in_channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(in_channels // reduction, in_channels, bias=False),
nn.Sigmoid()
)
def forward(self, x):
b, c, _, _ = x.size() # 2，128，20，20
y = self.avg_pool(x).view(b, c)
y = self.fc(y).view(b, c, 1, 1)
return x * y.expand_as(x)
class ConvSSM1(nn.Module):
def __init__(
self,
hidden_dim: int = 0,
drop_path: float = 0,
norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
attn_drop_rate: float = 0,
d_state: int = 16,
**kwargs,
):
super().__init__()
self.ln_1 = norm_layer(hidden_dim // 2)
self.self_attention = SS2D(d_model=hidden_dim // 2, dropout=attn_drop_rate, d_state=d_state, **kwargs)
self.drop_path = DropPath(drop_path)
self.frm = nn.Sequential(
nn.Conv2d(in_channels=hidden_dim // 2, out_channels=hidden_dim // 4, kernel_size=1, bias=False),
nn.BatchNorm2d(hidden_dim // 4),
nn.ReLU(inplace=True),
AttentionBlockWithMLP(hidden_dim // 4),
nn.Conv2d(hidden_dim // 4, hidden_dim // 2, kernel_size=1, bias=False),
nn.BatchNorm2d(hidden_dim // 2)
)
self.finalconv11 = nn.Conv2d(in_channels=hidden_dim, out_channels=hidden_dim, kernel_size=1, stride=1)
def forward(self, input: torch.Tensor):
input_left, input_right = input.chunk(2, dim=-1)
x = input_right + self.drop_path(self.self_attention(self.ln_1(input_right)))
input_left = input_left.permute(0, 3, 1, 2).contiguous()
input_left = self.frm(input_left)
x = x.permute(0, 3, 1, 2).contiguous()
output = torch.cat((input_left, x), dim=1)
output = self.finalconv11(output).permute(0, 2, 3, 1).contiguous()
return output + input
@MODELS.register_module()
class ConvSSMFE(nn.Module):
def __init__(self, drop_path,
attn_drop_rate, d_state,
in_channels,
widen_factor: float = 1.0,
):
super().__init__()
in_channels = [
int(channel * widen_factor) for channel in in_channels
],
self.convssm0 = ConvSSM1(hidden_dim=in_channels[0][0], drop_path=drop_path, norm_layer=nn.LayerNorm,
attn_drop_rate=attn_drop_rate, d_state=d_state)
self.convssm1 = ConvSSM1(hidden_dim=in_channels[0][1], drop_path=drop_path, norm_layer=nn.LayerNorm,
attn_drop_rate=attn_drop_rate, d_state=d_state)
self.convssm2 = ConvSSM1(hidden_dim=in_channels[0][2], drop_path=drop_path, norm_layer=nn.LayerNorm,
attn_drop_rate=attn_drop_rate, d_state=d_state)
def forward(self, x):
output = []
x1, x2, x3 = x
x1 = x1.permute(0, 2, 3, 1)
x2 = x2.permute(0, 2, 3, 1)
x3 = x3.permute(0, 2, 3, 1)
x1 = self.convssm0(x1)
x2 = self.convssm1(x2)
x3 = self.convssm2(x3)
x1 = x1.permute(0, 3, 1, 2)
x2 = x2.permute(0, 3, 1, 2)
x3 = x3.permute(0, 3, 1, 2)
output.append(x1)
output.append(x2)
output.append(x3)
return tuple(output)
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ------------------训练测试，反向传播报错-------------------------------------
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.nn import functional as F
# 定义虚拟数据集
class FakeDetectionDataset(Dataset):
def __init__(self, num_samples=100):
self.num_samples = num_samples
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
# 随机生成图像和边界框标签
image = torch.randn(20, 20, 512) # 假设图像通道数为1
boxes = torch.tensor([[100, 100, 200, 200]], dtype=torch.float) # 假设有一个边界框
labels = torch.tensor([1], dtype=torch.float).unsqueeze(0) # 假设边界框是目标类别1
return image, boxes, labels
# 定义简单的目标检测神经网络
class SimpleDetectionNet(nn.Module):
def __init__(self, num_classes=1):
super(SimpleDetectionNet, self).__init__()
self.conv_ssm1 = ConvSSM1(hidden_dim=512, drop_path=0.1, norm_layer=nn.LayerNorm, attn_drop_rate=0.1,
d_state=16)
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512, num_classes + 4) # 4为边界框的坐标和宽高
def forward(self, x):
x = self.conv_ssm1(x)
x = x.permute(0, 3, 1, 2)
x = self.pool(x) # 2 512 1 1
x = torch.flatten(x, 1) # 2 512
output = self.fc(x)
return output
# 定义自定义损失函数
def custom_detection_loss(pred, target_class, target_bbox):
pred_cls = torch.round(pred[:, -1:])
class_loss = F.cross_entropy(pred_cls, target_class)
bbox_loss = F.smooth_l1_loss(pred[:, :-2], target_bbox)
total_loss = class_loss + bbox_loss
return total_loss
# 实例化数据集和数据加载器
dataset = FakeDetectionDataset(num_samples=100)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)
# 实例化网络和设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = SimpleDetectionNet(num_classes=2).to(device) # 假设有两个类别，包括背景
# 定义优化器
optimizer = optim.Adam(net.parameters(), lr=0.001)
# 训练模型
num_epochs = 2
for epoch in range(num_epochs):
for images, boxes, labels in dataloader:
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
# 前向传播
outputs = net(images)
# 计算损失
loss = custom_detection_loss(outputs, labels[:, 0], boxes[:, 0])
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'Epoch [{epoch + 1}/{num_epochs}], Batch Loss: {loss.item():.4f}')
print("Training complete.")