sage.architecture.network.mscnn1d_att_resnet2d_cbam

Filename : legacy.py Description : Short description of the file

Created on 2025-11-06 12:57:18

__author__ = Narenraju Nagarajan __copyright__ = Copyright 2025, ProjectName __license__ = MIT Licence __version__ = 0.0.1 __maintainer__ = Narenraju Nagarajan __affiliation__ = N/A __email__ = N/A __status__ = [‘inProgress’, ‘Archived’, ‘inUsage’, ‘Debugging’]

GitHub Repository: NULL

Documentation: NULL

Classes

MSCNN1D_2DResNetCBAM	Multi-scale CNN backend + ResNet CBAM frontend for GW detection.
MSCNN1D_2DResNetCBAM_Heteroscedastic	Multi-scale CNN frontend + ResNet CBAM backend for GW detection.
ConsistencyOutput	Forward output of MSCNN1D_2DResNetCBAM_Consistency.
MSCNN1D_2DResNetCBAM_Consistency	Heteroscedastic detection network + multi-detector consistency heads.

Module Contents

class MSCNN1D_2DResNetCBAM(frontend_filters=32, frontend_kernel=64, backend_resnet_size=50, norm_type='groupnorm', dropout=0.0)

Bases: torch.nn.Module

Multi-scale CNN backend + ResNet CBAM frontend for GW detection.

Parameters:

• backend_filters – base filter size for ConvBlock backend

• backend_kernel – base kernel size for ConvBlock backend

• resnet_size – 18, 34, 50, 101, 152

• norm_type (str) – ‘batchnorm’, ‘layernorm’, ‘instancenorm’

• num_point_estimates – number of continuous parameters to predict

• frontend_filters (int)

• frontend_kernel (int)

• backend_resnet_size (int)

• dropout (float)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

num_detectors

norm

frontend

backend

avg_pool_1d

flatten

get_ranking_statistic

point_estimate_layers

forward(x)

x: Tensor of shape (batch, 2, signal_length) returns: raw, pred_prob, point_estimates (list of tensors)

class MSCNN1D_2DResNetCBAM_Heteroscedastic(frontend_filters=32, frontend_kernel=64, backend_resnet_size=50, norm_type='groupnorm', dropout=0.0)

Bases: torch.nn.Module

Multi-scale CNN frontend + ResNet CBAM backend for GW detection. Outputs:

• Ranking statistic (BCE)

• Point estimates (mean + log variance for heteroscedastic regression)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• frontend_filters (int)

• frontend_kernel (int)

• backend_resnet_size (int)

• norm_type (str)

• dropout (float)

num_detectors

norm

frontend

backend

avg_pool_1d

flatten

get_ranking_statistic

point_estimate_layers

forward(x)

x: Tensor of shape (B, num_detectors=2, signal_length) :returns: (B, 1)

point_estimates: (B, 2*num_pe)

Blocked format: [mu_0, mu_1, …, sraw_0, sraw_1, …] (all predicted means first, then all raw sigma params). The raw sigma params are mapped to a strictly-positive std via softplus inside BCEWithPEsigmaLoss (no exp(log_var) collapse). This matches the layout expected by BCEWithPEsigmaLoss and SageUncompiledValidation, which split at [:num_pe] / [num_pe:].

Return type:

ranking_statistic

class ConsistencyOutput

Bases: NamedTuple

Forward output of MSCNN1D_2DResNetCBAM_Consistency.

ranking_stat: torch.Tensor

point_estimates: torch.Tensor

mu_tc: torch.Tensor

sigma_tc: torch.Tensor

mu_mc: torch.Tensor

sigma_mc: torch.Tensor

s_tc: torch.Tensor

s_mc: torch.Tensor

class MSCNN1D_2DResNetCBAM_Consistency(t_grid, frontend_filters=32, frontend_kernel=64, backend_resnet_size=50, norm_type='groupnorm', head_hidden=128, dropout=0.0)

Bases: torch.nn.Module

Heteroscedastic detection network + multi-detector consistency heads.

Identical merged backbone (two frontends -> 2-D ResNet-CBAM -> ranking + heteroscedastic PE heads) to MSCNN1D_2DResNetCBAM_Heteroscedastic, plus a shared PerDetHead applied to each frontend output (pre-merge). The per-detector tc/mchirp estimates form an uncertainty-weighted consistency statistic (s_tc, s_mc); those, the per-detector sigmas and the attention entropies are concatenated into the ranking head input as a learned corroboration combiner (not a hard gate).

Parameters:

• t_grid (torch.Tensor, shape (L,)) – Physical time (seconds) of each multirate output sample, from output_time_grid(). Adaptive-averaged to the frontend’s time length to give the per-step t_position the tc soft-argmax sums over.

• frontend_filters (int) – As in the heteroscedastic model.

• frontend_kernel (int) – As in the heteroscedastic model.

• backend_resnet_size (int) – As in the heteroscedastic model.

• norm_type (str) – As in the heteroscedastic model.

• head_hidden (int) – Hidden width of the per-detector head MLPs.

• state (Initialize internal Module)

• ScriptModule. (shared by both nn.Module and)

• dropout (float)

num_detectors

norm

frontend

backend

avg_pool_1d

flatten

per_det_head

tc_scale

n_corr = 8

corr_gain_init = 0.25

corr_norm

get_ranking_statistic

point_estimate_layers

forward(x, return_embedding=False)

Returns a ConsistencyOutput. When return_embedding is True, returns (ConsistencyOutput, embedding) instead, where embedding is the per-detector attention-pooled frontend feature, shape (B, D, C) – used by the hard-noise miner as its QD diversity descriptor. Default False keeps the standard output untouched.

Parameters:

return_embedding (bool)