Skill v1.0.0
currentAutomated scan100/100version: "1.0.0" name: cv-bigru-slice-feature-aggregator description: Two-stage CT classifier where a 2D CNN dumps per-slice features once, then a bidirectional GRU runs over the slice sequence to produce both per-slice predictions (TimeDistributed head) and an exam-level prediction (avg+max pooled head) — turns expensive 3D CNN training into cheap sequence modeling
Overview
Volumetric CT classification with a true 3D CNN is expensive: every epoch you re-encode the same slices that haven't changed. The cheaper, often-better alternative is two-stage. Stage 1 trains a 2D CNN on slices, then dumps a fixed-size feature vector per slice once and freezes. Stage 2 trains a tiny bidirectional GRU over the per-slice feature sequence, with two heads: a TimeDistributed Linear that produces per-slice predictions and a cat(avg_pool, max_pool) Linear that produces the exam-level prediction. Adding the inter-slice Z-gap as an extra input feature gives the GRU spatial context. Stage 2 is so cheap you can sweep dozens of hyperparameters in the time stage 1 takes for one epoch.
Quick Start
import torchimport torch.nn as nnclass TimeDistributed(nn.Module):def __init__(self, layer): super().__init__(); self.layer = layerdef forward(self, x): # (B, T, F) -> (B, T, F_out)B, T, F = x.shapereturn self.layer(x.reshape(B * T, F)).reshape(B, T, -1)class SliceGRU(nn.Module):def __init__(self, n_feats, hidden=64, n_exam_targets=9):super().__init__()self.gru = nn.GRU(n_feats + 1, # +1 for inter-slice z-gaphidden,num_layers=2,bidirectional=True,batch_first=True,)self.image_head = TimeDistributed(nn.Linear(hidden * 2, 1))self.exam_head = nn.Linear(hidden * 2 * 2, n_exam_targets)def forward(self, slice_feats, z_gaps):x = torch.cat([slice_feats, z_gaps.unsqueeze(-1)], dim=2)h, _ = self.gru(x) # (B, T, 2H)per_slice = self.image_head(h)avg = h.mean(dim=1)mx, _ = h.max(dim=1)per_exam = self.exam_head(torch.cat([avg, mx], dim=1))return per_slice, per_exam
Workflow
- Train a 2D CNN end-to-end on per-slice classification (or load a pretrained backbone)
- For every CT series, run the 2D CNN once and dump
(num_slices, n_feats)to disk as a single.npy - Compute the inter-slice Z gap from
ImagePositionPatient[2]deltas; first slice gets gap = 0 - Train the GRU with both losses summed: per-slice BCE (with masking for padding) + per-exam BCE
- Use
cat(mean_pool, max_pool)for the exam-level head — single-pool is consistently worse - Keep the GRU tiny (hidden=64, 2 layers) — it's a sequence aggregator, not a feature extractor
Key Decisions
- Freeze stage 1 before dumping: any backbone update invalidates the feature cache; freeze + dump + train stage 2 is the right order.
- Bidirectional, not unidirectional: PE / lesion / nodule context is symmetric; left-right context matters as much as right-left.
- avg+max concat for exam head: max captures "worst slice", avg captures "overall burden"; they're complementary.
- z-gap as input feature: lets the GRU compensate for variable slice spacing across studies.
- Train both heads jointly: per-slice loss provides dense supervision the per-exam head couldn't learn alone.