Skip to content

get_net_builder

anomaly_match.utils.get_net_builder

Neural network model builder and registry.

TestCNN(num_classes=2, in_channels=3)

Bases: Module

Minimal CNN for fast testing. Not for production use.

Source code in anomaly_match/utils/get_net_builder.py 
56
57
58
59
60
61
62
63
def __init__(self, num_classes=2, in_channels=3):
    super().__init__()
    self._conv_stem = nn.Conv2d(in_channels, 8, 3, stride=2, padding=1)
    self.features = nn.Sequential(
        nn.ReLU(),
        nn.AdaptiveAvgPool2d(1),
    )
    self._fc = nn.Linear(8, num_classes)

forward(x)

Run the forward pass.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required

Returns:

Type Description
Tensor

Classification logits.
Source code in anomaly_match/utils/get_net_builder.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Run the forward pass.

    Args:
        x: Input tensor.

    Returns:
        Classification logits.
    """
    x = self._conv_stem(x)
    x = self.features(x)
    x = x.flatten(1)
    return self._fc(x)

get_net_builder(net_name, pretrained=False, in_channels=3)

Create a neural network builder function for the specified architecture.

This function returns a builder function that creates a neural network with the specified architecture when called with num_classes and in_channels parameters. Uses timm (pytorch-image-models) as the backend for all EfficientNet variants.

Parameters:

Name Type Description Default
net_name str

Name of the network architecture, supported values: - efficientnet-lite0 through efficientnet-lite4 - efficientnet-b0 through efficientnet-b7 - test-cnn (for testing only)

 required
pretrained bool

If True, loads pretrained weights. Default is False.

 False
in_channels int

Number of input channels. Default is 3.

 3

Returns:

Type Description
Callable

A function that builds the network when called with (num_classes, in_channels)
Source code in anomaly_match/utils/get_net_builder.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
def get_net_builder(net_name: str, pretrained: bool = False, in_channels: int = 3) -> Callable:
    """Create a neural network builder function for the specified architecture.

    This function returns a builder function that creates a neural network with the
    specified architecture when called with num_classes and in_channels parameters.
    Uses timm (pytorch-image-models) as the backend for all EfficientNet variants.

    Args:
        net_name: Name of the network architecture, supported values:
            - efficientnet-lite0 through efficientnet-lite4
            - efficientnet-b0 through efficientnet-b7
            - test-cnn (for testing only)
        pretrained: If True, loads pretrained weights. Default is False.
        in_channels: Number of input channels. Default is 3.

    Returns:
        A function that builds the network when called with (num_classes, in_channels)
    """
    if net_name == "test-cnn":
        logger.debug("Using test-cnn model (for testing only)")

        def build_test_cnn(num_classes, in_channels, pretrained=None):
            return TestCNN(num_classes=num_classes, in_channels=in_channels)

        return build_test_cnn

    timm_name, use_pretrained = _resolve_timm_name(net_name, pretrained)
    logger.debug(
        f"Using {'pretrained' if use_pretrained else 'non-pretrained'} {net_name} "
        f"(timm: {timm_name})"
    )

    def build_model(
        num_classes, in_channels, _timm_name=timm_name, _pretrained=use_pretrained, pretrained=None
    ):
        effective_pretrained = pretrained if pretrained is not None else _pretrained
        if effective_pretrained:
            try:
                model = timm.create_model(
                    _timm_name,
                    pretrained=True,
                    num_classes=num_classes,
                    in_chans=in_channels,
                    cache_dir=str(_PRETRAINED_CACHE_DIR),
                )
            except Exception:
                logger.warning(
                    f"Bundled pretrained weights not available (clone with git-lfs to avoid "
                    f"re-downloading). Downloading {_timm_name} from HuggingFace."
                )
                model = timm.create_model(
                    _timm_name,
                    pretrained=True,
                    num_classes=num_classes,
                    in_chans=in_channels,
                )
        else:
            model = timm.create_model(
                _timm_name,
                pretrained=False,
                num_classes=num_classes,
                in_chans=in_channels,
            )
        # timm initialises the classifier with TF-EfficientNet's
        # ``1/sqrt(fan_in + fan_out)`` scale (``_init_weight_goog``), tuned for a
        # 1000-class head. For AnomalyMatch's 2-class head that is ~num_classes×
        # too large (weight std ~0.4 vs ~0.016), so the fresh head is wildly
        # overconfident at init: nearly every unlabeled image clears FixMatch's
        # ``p_cutoff`` (0.95) with a random pseudo-label, burning a garbage
        # decision boundary into the backbone from the first step. Reset to
        # PyTorch's default ``1/sqrt(fan_in)`` Linear init (as the previous
        # efficientnet_lite_pytorch backbone used) so pseudo-labels ramp in only
        # as the model genuinely learns.
        classifier = model.get_classifier()
        if isinstance(classifier, nn.Linear):
            classifier.reset_parameters()
        return model

    return build_model