Untitled
Anonymous
plain_text
02/20/2026 4:11 PM
10.1 KB
10
Indexable
from models.model_utils import MeanDecodeNode
from models.neck_utils import *
import torch
class DetectionNeck(nn.Module):
def __init__(self, in_channels, channels=128, attention=False, step_mode='m', backend='cupy', args=None):
super().__init__()
self.epsilon = 1e-4
self.nz, self.numel = {}, {}
s3_c, s4_c, s5_c, s6_c, s7_c = in_channels
self.out_channels = [channels] * 5
self.T = args.T
self.attention = attention
# Calculate feature map sizes based on input resolution
# Using ACTUAL backbone downsampling ratios: /8, /16, /32, /64, /128
self.h, self.w = args.image_shape
s3_h, s3_w = self.h // 8, self.w // 8 # 720/8=90, 1280/8=160
s4_h, s4_w = self.h // 16, self.w // 16 # 720/16=45, 1280/16=80
s5_h, s5_w = self.h // 32, self.w // 32 # 720/32=22.5->22, 1280/32=40
s6_h, s6_w = self.h // 64, self.w // 64 # 720/64=11.25->11, 1280/64=20
s7_h, s7_w = self.h // 128, self.w // 128 # 720/128=5.625->5, 1280/128=10
# But actual sizes from backbone are slightly different due to conv/pooling
# s5: 23x40, s6: 12x20, s7: 6x10
# We need to use the ACTUAL input sizes for upsampling targets
# The upsampled s7 should match s6's actual size (12x20), not calculated (11x20)
# Actual sizes based on backbone output (derived from your debug output)
# For 720x1280: s3=90x160, s4=45x80, s5=23x40, s6=12x20, s7=6x10
s3_h, s3_w = self.h // 8, self.w // 8
s4_h, s4_w = (self.h // 8 + 1) // 2, (self.w // 8) // 2 # ceil division for odd
s5_h, s5_w = (s4_h + 1) // 2, s4_w // 2 # 23, 40
s6_h, s6_w = (s5_h + 1) // 2, s5_w // 2 # 12, 20
s7_h, s7_w = s6_h // 2, s6_w // 2 # 6, 10
# Spatial upsampling - use calculated dimensions
self.s7_s6_up = SpikingUpBlock(s7_c, s6_c, (s6_h, s6_w), step_mode=step_mode, backend=backend, args=args)
self.s6_s5_up = SpikingUpBlock(channels, s5_c, (s5_h, s5_w), step_mode=step_mode, backend=backend, args=args)
self.s5_s4_up = SpikingUpBlock(channels, s4_c, (s4_h, s4_w), step_mode=step_mode, backend=backend, args=args)
self.s4_s3_up = SpikingUpBlock(channels, s3_c, (s3_h, s3_w), step_mode=step_mode, backend=backend, args=args)
# Spatial downsampling
self.s3_s4_down = SpikingDownBlock(channels, channels, padding=(0, 0), step_mode=step_mode, backend=backend,
args=args)
self.s4_s5_down = SpikingDownBlock(channels, channels, padding=(1, 0), step_mode=step_mode, backend=backend,
args=args)
self.s5_s6_down = SpikingDownBlock(channels, channels, padding=(1, 0), step_mode=step_mode, backend=backend,
args=args)
self.s6_s7_down = SpikingDownBlock(channels, s7_c, padding=(0, 0), step_mode=step_mode, backend=backend,
args=args)
# Upsampling stage fusion module
self.s6_up_fusion = SpikingFusionBlock(s6_c, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
self.s5_up_fusion = SpikingFusionBlock(s5_c, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
self.s4_up_fusion = SpikingFusionBlock(s4_c, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
self.s3_up_fusion = SpikingFusionBlock(s3_c, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
# Downsampling stage fusion module
self.s4_down_fusion = SpikingFusionBlock(channels, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
self.s5_down_fusion = SpikingFusionBlock(channels, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
self.s6_down_fusion = SpikingFusionBlock(channels, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
self.s7_down_fusion = SpikingFusionBlock(s7_c, channels, attention=self.attention, step_mode=step_mode,
backend=backend, args=args)
# output neuron
if args.decode == 'spiking':
self.out_module = nn.ModuleList([
nn.Sequential(
neuron.LIFNode(step_mode='m', backend='cupy'),
MeanDecodeNode(T=self.T),
),
nn.Sequential(
neuron.LIFNode(step_mode='m', backend='cupy'),
MeanDecodeNode(T=self.T),
),
nn.Sequential(
neuron.LIFNode(step_mode='m', backend='cupy'),
MeanDecodeNode(T=self.T),
),
nn.Sequential(
neuron.LIFNode(step_mode='m', backend='cupy'),
MeanDecodeNode(T=self.T),
),
nn.Sequential(
neuron.LIFNode(step_mode='m', backend='cupy'),
MeanDecodeNode(T=self.T),
),
])
def forward(self, source_features):
assert len(source_features) == 5
output_features = []
T = source_features[0].shape[0]
assert T == self.T
# Obtain input features
s3_in = source_features[0]
s4_in = source_features[1]
s5_in = source_features[2]
s6_in = source_features[3]
s7_in = source_features[4]
# Upsampling fusion
s6_up = self.s6_up_fusion(s6_in + self.s7_s6_up(s7_in))
s5_up = self.s5_up_fusion(s5_in + self.s6_s5_up(s6_up))
s4_up = self.s4_up_fusion(s4_in + self.s5_s4_up(s5_up))
s3_out = self.s3_up_fusion(s3_in + self.s4_s3_up(s4_up))
# Downsampling fusion
s4_out = self.s4_down_fusion(s4_up + self.s3_s4_down(s3_out))
s5_out = self.s5_down_fusion(s5_up + self.s4_s5_down(s4_out))
s6_out = self.s6_down_fusion(s6_up + self.s5_s6_down(s5_out))
s7_out = self.s7_down_fusion(s7_in + self.s6_s7_down(s6_out))
# activation
output_features.append(self.out_module[0](s3_out))
output_features.append(self.out_module[1](s4_out))
output_features.append(self.out_module[2](s5_out))
output_features.append(self.out_module[3](s6_out))
output_features.append(self.out_module[4](s7_out))
assert len(output_features) == 5
return output_features
def add_hooks(self, instance):
def get_nz(name):
def hook(model, input, output):
self.nz[name] += torch.count_nonzero(output)
self.numel[name] += output.numel()
return hook
self.hooks = {}
for name, module in self.named_modules():
if isinstance(module, instance):
self.nz[name], self.numel[name] = 0, 0
self.hooks[name] = module.register_forward_hook(get_nz(name))
def reset_nz_numel(self):
for name, module in self.named_modules():
self.nz[name], self.numel[name] = 0, 0
def get_nz_numel(self):
return self.nz, self.numel
Editor is loading...
Leave a Comment