fix(architectures): follow layer_sizes for any ResNet num_blocks

deeptab/architectures/resnet.py

@@ -78,24 +78,25 @@ def __init__(
 
         self.initial_layer = nn.Linear(input_dim, self.hparams.layer_sizes[0])
 
+        # Block i maps layer_sizes[i] to layer_sizes[i + 1]. Once the width
+        # schedule is exhausted, later blocks stay at the final width, and the
+        # output layer reads whatever width the last block produces.
+        layer_sizes = self.hparams.layer_sizes
+        width = layer_sizes[0]
         self.blocks = nn.ModuleList()
         for i in range(self.hparams.num_blocks):
-            input_dim = self.hparams.layer_sizes[i]
-            output_dim = (
-                self.hparams.layer_sizes[i + 1]
-                if i + 1 < len(self.hparams.layer_sizes)
-                else self.hparams.layer_sizes[-1]
-            )
+            output_dim = layer_sizes[min(i + 1, len(layer_sizes) - 1)]
             block = ResidualBlock(
-                input_dim,
+                width,
                 output_dim,
                 self.hparams.activation,
                 self.hparams.norm,
                 self.hparams.dropout,
             )
             self.blocks.append(block)
+            width = output_dim
 
-        self.output_layer = nn.Linear(self.hparams.layer_sizes[-1], num_classes)
+        self.output_layer = nn.Linear(width, num_classes)
 
     def forward(self, *data):
         """Forward pass of the ResNet model.

deeptab/configs/models/resnet_config.py

@@ -15,13 +15,17 @@ class ResNetConfig(BaseModelConfig):
     activation : Callable, default=nn.SELU()
         Activation function for the ResNet layers.
     layer_sizes : list, default=[256, 128, 32]
-        Sizes of the layers in the ResNet.
+        Width schedule of the ResNet. The initial projection maps the input to
+        ``layer_sizes[0]``, and residual block ``i`` maps ``layer_sizes[i]`` to
+        ``layer_sizes[i + 1]``.
     dropout : float, default=0.5
         Dropout rate for regularization.
     norm : bool, default=False
         Whether to use normalization in the ResNet.
     num_blocks : int, default=3
-        Number of residual blocks in the ResNet.
+        Number of residual blocks in the ResNet. Blocks beyond the width
+        schedule keep the last width of ``layer_sizes``; with fewer blocks the
+        output layer reads the width of the last block.
     """
 
     # Override parent defaults

docs/model_zoo/stable/resnet.md

@@ -25,16 +25,16 @@ features -> optional embeddings -> initial Linear -> ResidualBlock x num_blocks
 
 ## Main Building Blocks
 
-| Component            | DeepTab implementation                    | Role                                        |
-| -------------------- | ----------------------------------------- | ------------------------------------------- |
-| Input representation | Raw concatenation or `EmbeddingLayer`     | Converts heterogeneous columns to a tensor. |
-| Initial projection   | `nn.Linear(input_dim, layer_sizes[0])`    | Sets hidden width.                          |
-| Residual body        | `ResidualBlock`                           | Learns transformations with skip paths.     |
-| Output layer         | `nn.Linear(layer_sizes[-1], num_classes)` | Produces task outputs.                      |
+| Component            | DeepTab implementation                            | Role                                        |
+| -------------------- | ------------------------------------------------- | ------------------------------------------- |
+| Input representation | Raw concatenation or `EmbeddingLayer`             | Converts heterogeneous columns to a tensor. |
+| Initial projection   | `nn.Linear(input_dim, layer_sizes[0])`            | Sets hidden width.                          |
+| Residual body        | `ResidualBlock`                                   | Learns transformations with skip paths.     |
+| Output layer         | `nn.Linear(width of the last block, num_classes)` | Produces task outputs.                      |
 
 ## Implementation Notes
 
-`num_blocks` controls how many residual blocks are instantiated. Each block uses `layer_sizes[i]` as input width and `layer_sizes[i + 1]` when available, otherwise the last width is reused. Keep `num_blocks` aligned with the length of `layer_sizes`; if `num_blocks` exceeds the number of transitions, later blocks stay at the final width.
+`num_blocks` controls how many residual blocks are instantiated. Block `i` maps `layer_sizes[i]` to `layer_sizes[i + 1]`. If `num_blocks` exceeds the number of transitions, later blocks stay at the final width; with fewer blocks, the output layer reads the width of the last block.
 
 ## Practical Config
 

tests/architectures/test_architecture_contracts.py

@@ -367,3 +367,31 @@ def test_model_defines_key_attributes(model_class):
     expected_attrs = ["returns_ensemble"]
     for attr in expected_attrs:
         assert hasattr(model, attr), f"{model_class.__name__} should define '{attr}'."
+
+
+@pytest.mark.parametrize(
+    "num_blocks, expected_widths",
+    [
+        (0, []),
+        (1, [(16, 8)]),
+        (2, [(16, 8), (8, 4)]),
+        (3, [(16, 8), (8, 4), (4, 4)]),
+        (4, [(16, 8), (8, 4), (4, 4), (4, 4)]),
+        (5, [(16, 8), (8, 4), (4, 4), (4, 4), (4, 4)]),
+    ],
+    ids=["0-blocks", "1-block", "2-blocks", "3-blocks", "4-blocks", "5-blocks"],
+)
+def test_resnet_blocks_follow_layer_sizes_for_any_num_blocks(num_blocks, expected_widths):
+    model_class = next(model_class for model_class in model_classes if model_class.__name__ == "ResNet")
+    config = get_model_config(model_class)
+    config.layer_sizes = [16, 8, 4]
+    config.num_blocks = num_blocks
+    model = model_class(
+        ({f"value_{index}": {"dimension": 1} for index in range(3)}, {}, {}), num_classes=2, config=config
+    )
+    widths = [(block.linear1.in_features, block.linear1.out_features) for block in model.blocks]
+    assert widths == expected_widths
+    assert model.output_layer.in_features == (expected_widths[-1][1] if expected_widths else 16)
+    output = model([torch.randn(5, 1) for _ in range(3)], [], [])
+    assert output.shape == (5, 2)
+    output.sum().backward()