SS2D反向傳播問題記錄【未解決】

使用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

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 DropPathfrom mmyolo.registry import MODELS  # newtry:from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, selective_scan_ref
except Exception as e:# print(e, flush=True)passtry:"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_cudaclass SelectiveScanMamba(torch.autograd.Function):# comment all checks if inside cross_selective_scan@staticmethod@torch.cuda.amp.custom_fwddef 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_softplusout, 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_bwddef backward(ctx, dout, *args):u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensorsif 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_fwddef 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_softplusif 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_bwddef backward(ctx, dout, *args):u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensorsif 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_fwddef 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_softplusctx.backnrows = backnrowsx = deltaout = uctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)return out@staticmethod@torch.cuda.amp.custom_bwddef backward(ctx, dout, *args):u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensorsif 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):@staticmethoddef forward(ctx, x: torch.Tensor):B, C, H, W = x.shapectx.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@staticmethoddef backward(ctx, ys: torch.Tensor):# out: (b, k, d, l)B, C, H, W = ctx.shapeL = 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_resdef 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) % Wexpanded_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)) % Wexpanded_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_shapeshift = torch.arange(H, device=tensor_flat.device).unsqueeze(1)index = (shift + torch.arange(W, device=tensor_flat.device)) % Wexpanded_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_tensordef antidiagonal_scatter(tensor_flat, original_shape):B, C, H, W = original_shapeshift = torch.arange(H, device=tensor_flat.device).unsqueeze(1)index = (torch.arange(W, device=tensor_flat.device) - shift) % Wexpanded_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_tensorclass CrossMerge(torch.autograd.Function):@staticmethoddef forward(ctx, ys: torch.Tensor):B, K, D, H, W = ys.shapectx.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_dareturn y_res.view(B, D, -1)# return yclass CrossScan_Ab_1direction(torch.autograd.Function):@staticmethoddef forward(ctx, x: torch.Tensor):B, C, H, W = x.shapectx.shape = (B, C, H, W)xs = x.view(B, 1, C, H * W).repeat(1, 4, 1, 1).contiguous()return xs@staticmethoddef backward(ctx, ys: torch.Tensor):# out: (b, k, d, l)B, C, H, W = ctx.shapey = ys.sum(dim=1).view(B, C, H, W)return yclass CrossMerge_Ab_1direction(torch.autograd.Function):@staticmethoddef forward(ctx, ys: torch.Tensor):B, K, D, H, W = ys.shapectx.shape = (H, W)y = ys.sum(dim=1).view(B, D, H * W)return y@staticmethoddef backward(ctx, x: torch.Tensor):# B, D, L = x.shape# out: (b, k, d, l)H, W = ctx.shapeB, C, L = x.shapexs = x.view(B, 1, C, L).repeat(1, 4, 1, 1).contiguous().view(B, 4, C, H, W)return xsclass CrossScan_Ab_2direction(torch.autograd.Function):@staticmethoddef forward(ctx, x: torch.Tensor):B, C, H, W = x.shapectx.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@staticmethoddef backward(ctx, ys: torch.Tensor):# out: (b, k, d, l)B, C, H, W = ctx.shapeL = H * Wys = 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):@staticmethoddef forward(ctx, ys: torch.Tensor):B, K, D, H, W = ys.shapectx.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@staticmethoddef backward(ctx, x: torch.Tensor):# B, D, L = x.shape# out: (b, k, d, l)H, W = ctx.shapeB, C, L = x.shapexs = x.new_empty((B, 4, C, L))xs[:, 0] = xxs[:, 1] = xxs[:, 2:4] = torch.flip(xs[:, 0:2], dims=[-1])xs = xs.view(B, 4, C, H, W)return xsdef 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 dtypeforce_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.shapeD, N = A_logs.shapeK, D, R = dt_projs_weight.shapeL = H * Wif nrows == 0:if D % 4 == 0:nrows = 4elif D % 3 == 0:nrows = 3elif D % 2 == 0:nrows = 2else:nrows = 1if backnrows == 0:if D % 4 == 0:backnrows = 4elif D % 3 == 0:backnrows = 3elif D % 2 == 0:backnrows = 2else:backnrows = 1def 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):@staticmethoddef forward(ctx, ys: torch.Tensor):B, K, D, H, W = ys.shapectx.shape = (H, W)y = ys.sum(dim=1).view(B, D, H * W)return y@staticmethoddef backward(ctx, x: torch.Tensor):# B, D, L = x.shape# out: (b, k, d, l)H, W = ctx.shapeB, C, L = x.shapexs = x.view(B, 1, C, L).repeat(1, 4, 1, 1).contiguous().view(B, 4, C, H, W)return xsclass 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 convconv_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_rankself.d_conv = d_conv# tags for forward_type ==============================def checkpostfix(tag, value):ret = value[-len(tag):] == tagif ret:value = value[:-len(tag)]return ret, valueself.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, Dsself.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) < 1self.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, Dsself.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) < 1self.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)))@staticmethoddef 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 initializationdt_init_std = dt_rank ** -0.5 * dt_scaleif 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_maxdt = 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/72759inv_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 = Truereturn dt_proj@staticmethoddef A_log_init(d_state, d_inner, copies=-1, device=None, merge=True):# S4D real initializationA = 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 fp32if 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 = Truereturn A_log@staticmethoddef D_init(d_inner, copies=-1, device=None, merge=True):# D "skip" parameterD = 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 fp32D._no_weight_decay = Truereturn D# only used to run previous versiondef 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.shapeD, N = self.A_logs.shapeK, D, R = self.dt_projs_weight.shapeL = H * Wx_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.floatinv_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_yy = 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 xdef forward(self, x: torch.Tensor, **kwargs):with_dconv = (self.d_conv > 1)  # Truex = 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 outclass 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，20y = 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 = xx1 = 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 torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import DataLoader, Datasetfrom torch.nn import functional as F# 定義虛擬數據集class FakeDetectionDataset(Dataset):def __init__(self, num_samples=100):self.num_samples = num_samplesdef __len__(self):return self.num_samplesdef __getitem__(self, idx):# 隨機生成圖像和邊界框標簽image = torch.randn(20, 20, 512)  # 假設圖像通道數為1boxes = torch.tensor([[100, 100, 200, 200]], dtype=torch.float)  # 假設有一個邊界框labels = torch.tensor([1], dtype=torch.float).unsqueeze(0)  # 假設邊界框是目標類別1return 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 1x = torch.flatten(x, 1)  # 2 512output = 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_lossreturn 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 = 2for 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.")

本文來自互聯網用戶投稿，該文觀點僅代表作者本人，不代表本站立場。本站僅提供信息存儲空間服務，不擁有所有權，不承擔相關法律責任。
如若轉載，請注明出處：http://www.pswp.cn/web/24162.shtml
繁體地址，請注明出處：http://hk.pswp.cn/web/24162.shtml
英文地址，請注明出處：http://en.pswp.cn/web/24162.shtml

如若內容造成侵權/違法違規/事實不符，請聯系多彩編程網進行投訴反饋email:809451989@qq.com，一經查實，立即刪除！

相關文章

AI大模型日報#0607：10家國產大模型、GPT-4o挑戰高考作文 | OpenAI公開破解GPT-4新方法

AI大模型日報#0607：10家國產大模型、GPT-4o挑戰高考作文 | OpenAI公開破解GPT-4新方法

導讀&#xff1a;AI大模型日報&#xff0c;爬蟲LLM自動生成&#xff0c;一文覽盡每日AI大模型要點資訊&#xff01;目前采用“文心一言”&#xff08;ERNIE 4.0&#xff09;、“零一萬物”&#xff08;Yi-Large&#xff09;生成了今日要點以及每條資訊的摘要。歡迎閱讀&#xf…
閱讀更多...
TS 系列：使用元祖生成聯合類型

TS 系列：使用元祖生成聯合類型

需求&#xff1a;有這么個需求&#xff0c;我們有兩個數組&#xff0c;一個記錄撲克牌花色&#xff0c;一個記錄撲克牌點數&#xff0c;需要有一個函數&#xff0c;傳遞兩個值&#xff0c;根據傳遞的值生成撲克牌&#xff0c;需要我們定義參數的類型檢查。 思路&#xff1a;肯…
閱讀更多...
2024速通python之python高階技巧

2024速通python之python高階技巧

文章目錄 一、閉包1.什么是閉包2.優缺點3.nonlocal關鍵字 二、裝飾器1.什么是裝飾器2.舉例3.傳統方式4.裝飾器方式5.語法糖寫法 三、多線程1.線程參數2.多線程編程 四、網絡編程1.Socket服務端編程2.Socket客戶端編程 「章節總覽」 ??????【2024速通python之python基礎…
閱讀更多...
超過20W個高質量組件的開源PCB庫

超過20W個高質量組件的開源PCB庫

項目介紹 Celestial Altium Library是由Altium行業專家Mark Harris創建的一個龐大的免費開源數據庫庫&#xff0c;專為Altium Designer而設計&#xff0c;庫中包含超過20萬個優質組件 . 特點 高質量數據&#xff1a;Celestial Altium Library注重數據的質量&#xff0c;用戶可…
閱讀更多...
Day14：響應式網頁

Day14：響應式網頁

通過媒體查詢、Bootstrap 框架完成騰訊全端網頁響應式布局。 一、響應式布局方案 1、什么是響應式布局 它的主要特點是能夠使網頁根據不同的設備屏幕尺寸&#xff08;如桌面電腦、平板電腦、手機等&#xff09;和分辨率自動調整布局和顯示效果&#xff0c;以提供最佳的用戶體…
閱讀更多...
python的plt.axis()、plt.xlim() 和 plt.ylim()函數

python的plt.axis()、plt.xlim() 和 plt.ylim()函數

坐標軸相關設置1、plt.axis():坐標軸設置 函數 plt.axis(*v, **kwargs) 主要用于設置坐標軸的屬性，返回值為當前的坐標軸范圍 [xmin, xmax, ymin, ymax]，幾種調用方式如下： 調用方式 說明 axis() 返回當前的坐標軸范圍 [xmin, xmax, ymi…
閱讀更多...
pytorch構建模型訓練數據集

pytorch構建模型訓練數據集

pytorch構建模型訓練數據集 pytorch構建模型訓練數據集1.AlexNet:1.1.導入必要的庫&#xff1a;1.2.數據預處理和增強&#xff1a;1.3.加載數據集&#xff1a;1.4.劃分測試集和訓練集&#xff1a;1.5.創建數據加載器&#xff1a;1.6.加載AlexNet模型&#xff1a;1.7.修改模型以…
閱讀更多...
qt qDebug兼容LOGE

qt qDebug兼容LOGE

目錄 普通qDebug用法 qt qDebug兼容LOGE 模板參數2實現 qDebug 實現LOGE一樣的用法，這樣Android和qt同時支持LOGE打印日志 普通qDebug用法 #include <QApplication> #include <QDebug>int main(int argc, char *argv[]) {QApplication app(argc, argv);int ret…
閱讀更多...
帝國cms接入騰訊云人臉識別認證代碼

帝國cms接入騰訊云人臉識別認證代碼

利用帝國cms在做一些會員系統的時候&#xff0c;需要做人臉識別認證&#xff0c;之前接入了某api接口&#xff0c;發現身份證識別率真的低&#xff0c;還好充值的少&#xff0c;否則要出問題&#xff0c;后來發現會員注冊率降低了不少&#xff0c;最終還是決定使用騰訊云的人臉…
閱讀更多...
三十五篇：數字化轉型的引擎：賦能企業的ERP系統全景

三十五篇：數字化轉型的引擎：賦能企業的ERP系統全景

數字化轉型的引擎&#xff1a;賦能企業的ERP系統全景 1. 引言 在這個快速變化的數字時代&#xff0c;現代企業面臨著前所未有的挑戰和機遇。企業資源計劃&#xff08;ERP&#xff09;系統&#xff0c;作為數字化轉型的核心&#xff0c;扮演著至關重要的角色。它不僅是企業運營…
閱讀更多...
第二十六章HTML與CSS書寫規范

第二十六章HTML與CSS書寫規范

1.HTML書寫規范 1.文檔類型聲明及編碼 統一為html5 聲明類型。編碼統一為utf-8。 2.頁面tdk TDK是一個縮寫&#xff0c;其中“T”表示為網頁定義標題&#xff0c;“D”表示為網頁定義描述 description&#xff0c;“K”表示為搜索引擎定義關鍵詞keywords。 1、<title&g…
閱讀更多...
NFT Insider #133：蘇富比將拍賣BAYC系列NFT，Taiko創世NFT系列已上線

NFT Insider #133：蘇富比將拍賣BAYC系列NFT，Taiko創世NFT系列已上線

引言&#xff1a;NFT Insider由NFT收藏組織WHALE Members &#xff08;https://twitter.com/WHALEMembers&#xff09;、BeepCrypto &#xff08;https://twitter.com/beep_crypto&#xff09;聯合出品&#xff0c;濃縮每周NFT新聞&#xff0c;為大家帶來關于NFT最全面、最新鮮…
閱讀更多...
立創·天空星開發板-GD32F407VE-EXTI

立創·天空星開發板-GD32F407VE-EXTI

本文以 立創天空星開發板-GD32F407VET6-青春版 作為學習的板子&#xff0c;記錄學習筆記。 立創天空星開發板-GD32F407VE-EXTI 中斷硬件觸發中斷示例軟件觸發中斷示例 中斷 中斷分為內部中斷和外部中斷 外部中斷是由外部設備&#xff08;如按鍵、傳感器、通信接口等&#xff09…
閱讀更多...
rollup.js（入門篇）

rollup.js（入門篇）

前沿 Rollup 是一個用于 JavaScript 的模塊打包工具&#xff0c;它將小的代碼片段編譯成更大、更復雜的代碼&#xff0c;例如庫或應用程序。它使用 JavaScript 的 ES6 版本中包含的新標準化代碼模塊格式&#xff0c;而不是以前的 CommonJS 和 AMD 等特殊解決方案。ES 模塊允許…
閱讀更多...
如何在Python中處理時間和日期

如何在Python中處理時間和日期

在Python中處理時間和日期&#xff0c;我們通常使用內置的datetime模塊。這個模塊提供了類來操作日期和時間&#xff0c;包括date、time、datetime、timedelta以及tzinfo。以下是datetime模塊的一些基本用法&#xff1a; 1. 獲取當前日期和時間 python復制代碼 from datetime…
閱讀更多...
Linux編程for、while循環if判斷以及case語句用法

Linux編程for、while循環if判斷以及case語句用法

簡介 語法描述if條件語句if else條件判斷語句if else-if else多條件判斷語句for循環執行命令while循環執行命令until直到條件為真時停止循環case ... esac多選擇語句break跳出循環continue跳出當前循環 1. for 循環 for語句&#xff0c;定量循環&#xff0c;可以遍歷一個列表…
閱讀更多...
一文讀懂AI時代GPU的內存新寵-HBM

一文讀懂AI時代GPU的內存新寵-HBM

一文讀懂GPU最強輔助&#xff1a;HBM HBM&#xff0c;即高帶寬內存&#xff0c;是一項領先的3D堆疊DRAM技術&#xff0c;專為高性能計算和圖形處理單元&#xff08;GPU&#xff09;設計&#xff0c;滿足其對內存帶寬和容量的極致需求。該技術由AMD與海力士攜手研發&#xff0c;…
閱讀更多...
Apache安裝教程

Apache安裝教程

目錄 一、Apache知識點 Apache服務簡介 Apache下載網址 Apache的主要特點 二、Apache服務的搭建 1. 關閉防火墻 2. 安裝依賴環境以及編譯工具 3. 將apache安裝包拖入xshell內 4. 解壓壓縮包 5. 進入httpd主包 6. 指定安裝路徑&#xff0c;啟用字符集支持等 7. 優化執…
閱讀更多...
重慶工商大學社會工作專業試題及答案，分享幾個實用搜題和學習工具 #媒體#學習方法#知識分享

重慶工商大學社會工作專業試題及答案，分享幾個實用搜題和學習工具 #媒體#學習方法#知識分享

搜題軟件一般都是通過識別題目內容搜索出問題的答案&#xff0c;當識別內容不正確或搜索不到答案時&#xff0c;又得重新到其他軟件進行重復的操作&#xff0c;很是麻煩。所以我們可以使用專業的識別工具&#xff0c;對題目內容進行識別&#xff0c;然后把提取出來的內容單獨保…
閱讀更多...
流量錄制學習

流量錄制學習

AREX Cloud | AREX (arextest.com) 流量錄制學習&#xff0c;比vivo的moonbox要好用
閱讀更多...
最新文章

Copyright @ 2022~2023