diff --git a/deeptab/architectures/resnet.py b/deeptab/architectures/resnet.py index 0556b7af..a75c45a7 100644 --- a/deeptab/architectures/resnet.py +++ b/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. diff --git a/deeptab/configs/models/resnet_config.py b/deeptab/configs/models/resnet_config.py index b5c59e11..dff8004e 100644 --- a/deeptab/configs/models/resnet_config.py +++ b/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 diff --git a/docs/model_zoo/stable/resnet.md b/docs/model_zoo/stable/resnet.md index 410bd4e2..196aa259 100644 --- a/docs/model_zoo/stable/resnet.md +++ b/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 diff --git a/tests/architectures/test_architecture_contracts.py b/tests/architectures/test_architecture_contracts.py index ec9090b5..6142f139 100644 --- a/tests/architectures/test_architecture_contracts.py +++ b/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()