""" DenseNet3D V5 - Dose-Preserving Architecture. Key changes from V4: 1. BatchNorm3d to preserve absolute dose magnitude across samples 2. Less pooling - only 2 transitions instead of 3 3. Anisotropic pooling - pools XY more than Z (matches CT spacing) 4. Same API as V4 for drop-in replacement """ import torch from torch import nn from typing import Optional, Tuple, Union class DenseBlockV5(nn.Module): """DenseBlock with BatchNorm for dose-magnitude preservation.""" def __init__( self, num_convs: int, in_channels: int, growth_rate: int, drop_rate: float = 0.0, kernel_size: int = 3, ): super().__init__() self.layers = nn.ModuleList() self.drop_rate = drop_rate for i in range(num_convs): current_channels = in_channels + i * growth_rate # Pre-activation: BatchNorm -> SiLU -> Conv self.layers.append( nn.Sequential( nn.BatchNorm3d(current_channels), nn.SiLU(inplace=True), nn.Conv3d( current_channels, growth_rate, kernel_size=kernel_size, padding=kernel_size // 2, ), ) ) def forward(self, x: torch.Tensor) -> torch.Tensor: for block in self.layers: y = block(x) if self.drop_rate > 0 and self.training: y = nn.functional.dropout(y, p=self.drop_rate, training=True) x = torch.cat((x, y), dim=1) return x class BlurPool3d(nn.Module): """Anti-aliased downsampling: Gaussian blur + strided subsampling. Uses a fixed binomial-3 ([1,2,1]) kernel to low-pass filter before subsampling, preventing checkerboard artifacts in gradients. Based on Zhang (2019), "Making Convolutional Networks Shift-Invariant Again". """ def __init__(self, channels: int, stride: Tuple[int, ...] = (1, 2, 2)): super().__init__() self.channels = channels self.stride = stride self.padding = 1 # for 3x3x3 kernel # Binomial-3 tent kernel — gives uniform gradient scaling with stride 2 a = torch.tensor([1.0, 2.0, 1.0]) filt = a[:, None, None] * a[None, :, None] * a[None, None, :] filt = filt / filt.sum() filt = filt[None, None].repeat(channels, 1, 1, 1, 1) self.register_buffer("filt", filt) def forward(self, x: torch.Tensor) -> torch.Tensor: return nn.functional.conv3d( nn.functional.pad(x, [self.padding] * 6, mode="reflect"), self.filt, stride=self.stride, groups=self.channels, ) class DenseNet3DSmooth(nn.Module): """ DenseNet3D V5 - Same API as V4 but with dose-preserving changes. Changes from V4: - BatchNorm3d to preserve absolute dose magnitude - Anisotropic pooling (less Z pooling) - 3 dense blocks with 2 transitions (less aggressive pooling) - Anti-aliased downsampling (BlurPool) for smooth dose gradients """ def __init__( self, *, in_channels_ct: int = 1, in_channels_dose: int = 1, in_channels_mask: int = 1, stem_out_channels_total: int = 96, growth_rate: int = 32, arch: Tuple[int, int, int] = (4, 4, 4), block_ct_mask_input_grads: bool = True, drop_rate: float = 0.0, use_skip_connections: bool = True, batch_norm_dose: bool = True, ) -> None: super().__init__() if in_channels_dose < 1: raise ValueError("in_channels_dose must be >= 1") use_ct = in_channels_ct > 0 use_mask = in_channels_mask > 0 n_stems = 1 + int(use_ct) + int(use_mask) if stem_out_channels_total < n_stems: raise ValueError(f"stem_out_channels_total must be >= {n_stems}") base = stem_out_channels_total // n_stems rem = stem_out_channels_total - base * n_stems stem_dose_out = base + (1 if rem > 0 else 0) stem_ct_out = (base + (1 if rem > 1 else 0)) if use_ct else 0 stem_mask_out = base if use_mask else 0 self.in_channels_ct = in_channels_ct self.in_channels_dose = in_channels_dose self.in_channels_mask = in_channels_mask self.block_ct_mask_input_grads = block_ct_mask_input_grads self.use_skip_connections = use_skip_connections self.batch_norm_dose = batch_norm_dose # --- Anisotropic Stem with BatchNorm + anti-aliased downsampling --- def make_stem(in_ch: int, out_ch: int, use_batch_norm: bool = True) -> nn.Sequential: layers = [ nn.Conv3d( in_ch, out_ch, kernel_size=(3, 7, 7), stride=1, # No stride — avoids checkerboard gradients padding=(1, 3, 3), ), ] if use_batch_norm: layers.append(nn.BatchNorm3d(out_ch)) layers.extend( [ nn.SiLU(inplace=True), # Anti-aliased 4x downsampling in H,W (two 2x stages) BlurPool3d(out_ch, stride=(1, 2, 2)), BlurPool3d(out_ch, stride=(1, 2, 2)), ] ) return nn.Sequential(*layers) self.stem_ct = make_stem(in_channels_ct, stem_ct_out) if use_ct else None self.stem_dose = make_stem(in_channels_dose, stem_dose_out, use_batch_norm=batch_norm_dose) self.stem_mask = make_stem(in_channels_mask, stem_mask_out) if use_mask else None init_channels = ( stem_dose_out + (stem_ct_out if use_ct else 0) + (stem_mask_out if use_mask else 0) ) # Anti-aliased transition: BlurPool instead of AvgPool for smooth gradients def make_transition(in_ch: int, out_ch: int, pool_z: bool = False) -> nn.Sequential: stride = (2, 2, 2) if pool_z else (1, 2, 2) return nn.Sequential( nn.BatchNorm3d(in_ch), nn.SiLU(inplace=True), nn.Conv3d(in_ch, out_ch, kernel_size=1), BlurPool3d(out_ch, stride=stride), ) # Build dense blocks and transitions (3 blocks, 2 transitions) self.dense_blocks = nn.ModuleList() self.transitions = nn.ModuleList() self.skip_projections = nn.ModuleList() if use_skip_connections else None out_channels = init_channels self._block_output_channels = [init_channels] # Pool schedule: first transition XY only, second transition includes Z pool_z_schedule = [False, True] for i, num_convs in enumerate(arch): self.dense_blocks.append( DenseBlockV5(num_convs, out_channels, growth_rate, drop_rate=drop_rate) ) out_channels += num_convs * growth_rate if i != len(arch) - 1: new_out = out_channels // 2 pool_z = pool_z_schedule[i] if i < len(pool_z_schedule) else True self.transitions.append(make_transition(out_channels, new_out, pool_z=pool_z)) if use_skip_connections: skip_stride = (2, 2, 2) if pool_z else (1, 2, 2) self.skip_projections.append( nn.Sequential( nn.Conv3d(self._block_output_channels[-1], new_out, kernel_size=1), BlurPool3d(new_out, stride=skip_stride), ) ) out_channels = new_out self._block_output_channels.append(out_channels) self._final_channels = out_channels # Final head self.final_norm = nn.BatchNorm3d(out_channels) self.final_act = nn.SiLU(inplace=True) self.global_pool = nn.AdaptiveAvgPool3d((1, 1, 1)) self.classifier = nn.Linear(out_channels, 1) self.apply(self._custom_init_weights) self._dose_input_ref: Optional[torch.Tensor] = None def forward( self, x_or_dose: Union[torch.Tensor, None], mask: Optional[torch.Tensor] = None, ct: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Forward pass - same API as V4.""" if mask is None and ct is None: x = x_or_dose if x is None: raise ValueError("Input tensor is None") b, c, d, h, w = x.shape c_ct = self.in_channels_ct c_dose = self.in_channels_dose c_mask = self.in_channels_mask expected = c_ct + c_dose + c_mask if c != expected: raise ValueError(f"Expected {expected} channels, got {c}") off = 0 x_ct = x[:, off : off + c_ct] if c_ct > 0 else None off += c_ct x_dose = x[:, off : off + c_dose] off += c_dose x_mask = x[:, off : off + c_mask] if c_mask > 0 else None else: x_dose = x_or_dose x_mask = mask x_ct = ct if x_dose is None: raise ValueError("Dose tensor must be provided") self._dose_input_ref = x_dose if self.block_ct_mask_input_grads: if x_ct is not None: x_ct = x_ct.detach() if x_mask is not None: x_mask = x_mask.detach() parts = [] if self.stem_ct is not None: if x_ct is None: raise ValueError("CT input missing") parts.append(self.stem_ct(x_ct)) parts.append(self.stem_dose(x_dose)) if self.stem_mask is not None: if x_mask is None: raise ValueError("Mask input missing") parts.append(self.stem_mask(x_mask)) x = torch.cat(parts, dim=1) skip_features = [x] if self.use_skip_connections else None for i, dense_block in enumerate(self.dense_blocks): x = dense_block(x) if i < len(self.transitions): x = self.transitions[i](x) if self.use_skip_connections and i < len(self.skip_projections): skip = self.skip_projections[i](skip_features[-1]) x = x + skip skip_features.append(x) x = self.final_norm(x) x = self.final_act(x) x = self.global_pool(x) x = x.view(x.size(0), -1) out = self.classifier(x) return out @staticmethod def _custom_init_weights(m: nn.Module) -> None: if isinstance(m, (nn.Conv3d, nn.Linear)): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm3d): if m.weight is not None: nn.init.ones_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) def enable_dose_gradient_mode(self) -> None: self.eval() self.block_ct_mask_input_grads = True