poolmult

timm's NormMlpClassifierHead builds an adaptive pool that accepts the documented catavgmax pool type, which concatenates average and max pooling and so doubles the channel dimension, but sizes its norm and classifier to the un-doubled channel count. So global_pool='catavgmax' feeds a 2x-channel tensor into a norm sized for 1x, and the forward raises a RuntimeError for every model that uses this head (ConvNeXt, ConvNeXtV2, MaxViT, CoAtNet). timm's other two classifier heads handle the doubling correctly, which is what makes this an oversight rather than an intended restriction.

The code

timm/layers/classifier.py, current main (identical in 1.0.28):

# NormMlpClassifierHead.__init__
self.global_pool = SelectAdaptivePool2d(pool_type=pool_type)   # accepts 'catavgmax'
self.norm = norm_layer(in_features, **dd)                       # sized in_features, no feat_mult

SelectAdaptivePool2d.feat_mult() returns 2 for catavgmax, so the pooled tensor has in_features * 2 channels. The two sibling heads in the same file account for this:

num_pooled_features = num_features * global_pool.feat_mult()        # create_classifier, line 33
num_pooled_features = self.in_features * self.global_pool.feat_mult()  # ClassifierHead, line 129
assert pool_type in ('', 'avg', 'max', 'avgmax')                   # ClNormMlpClassifierHead, line 254

ClassifierHead sizes its classifier by feat_mult, and ClNormMlpClassifierHead refuses the doubling pool types with an assert. NormMlpClassifierHead does neither: it neither multiplies by feat_mult nor guards the pool type, so it is the only one of the three that mis-sizes.

What it does

The channel counts show the mismatch directly:

pool=       avg: pooled channels = 320, ClassifierHead consumer = 320, NormMlp consumer = 320
pool=    avgmax: pooled channels = 320, ClassifierHead consumer = 320, NormMlp consumer = 320
pool= catavgmax: pooled channels = 640, ClassifierHead consumer = 640, NormMlp consumer = 320

and the real timm heads confirm the consequence:

NormMlpClassifierHead catavgmax: {'ok': False, 'error': 'RuntimeError'}
NormMlpClassifierHead avgmax (base case): {'ok': True, 'shape': (2, 10)}
ClassifierHead catavgmax (sibling): {'ok': True, 'shape': (2, 10)}

catavgmax produces 640 channels while the norm is sized for 320, so the forward raises RuntimeError: Given normalized_shape=[320] ... got input of size [2, 1, 1, 640]. The same crash happens through the public API, timm.create_model('convnext_atto', global_pool='catavgmax'). The base case confirms the harness: avgmax, the pool type that does not double channels, works on the exact same head, and the sibling ClassifierHead accepts catavgmax because it sizes with feat_mult. So the defect is isolated to the missing feat_mult in NormMlpClassifierHead, not to catavgmax itself.

Scope

The bug is on current main and is reached through the standard public APIs create_model(..., global_pool='catavgmax') and reset_classifier(..., global_pool='catavgmax'), which are the documented ways to select a pooling mode. It affects the ConvNeXt, ConvNeXtV2, MaxViT, and CoAtNet families that use this head. catavgmax is one of timm's four documented pooling options, and the fix is to size the norm and fc by in_features * self.global_pool.feat_mult() (matching ClassifierHead) or to add the ClNormMlpClassifierHead guard.

Layout

• excerpt.py: the NormMlpClassifierHead sizing and the two siblings quoted with line numbers and the three flags (Apache-2.0).
• head.py: the pooled-channel count and each head's consumer size, showing the catavgmax mismatch without a model.
• consequence.py: the real timm NormMlpClassifierHead crashing on catavgmax, the avgmax base case, and the sibling ClassifierHead working.
• test_poolmult.py: catavgmax doubles channels, ClassifierHead and the guarded head are consistent, NormMlpClassifierHead is not, and the real timm heads crash / work as described.

Reproduce

python head.py
python consequence.py
python test_poolmult.py

The sizing is quoted from current main; the crash is produced by the real timm heads. The fix is to size NormMlpClassifierHead by in_features * feat_mult() like its sibling.