Skip to content

fix: model_convnext_small_22k1k defaults to 1000 classes - #376

Merged
cregouby merged 2 commits into
mlverse:mainfrom
sebffischer:fix/convnext-small-22k1k-num-classes
Aug 1, 2026
Merged

fix: model_convnext_small_22k1k defaults to 1000 classes#376
cregouby merged 2 commits into
mlverse:mainfrom
sebffischer:fix/convnext-small-22k1k-num-classes

Conversation

@sebffischer

Copy link
Copy Markdown
Contributor

convnext_small_22k1k is pretrained on Imagenet 22k and fine-tuned on Imagenet 1k, so the published weights have a 1000-class head. The formal default was num_classes = 21841, taken from the other _22k models, so model_convnext_small_22k1k(pretrained = TRUE) built a 21841-class head and then failed to load the state dict with

The size of tensor a (21841) must match the size of tensor b (1000)
at non-singleton dimension 0

Verified against the published checkpoint: its head.weight is 1000x768, whereas convnext_tiny_22k really is 21841x768.

The caller could not work around this by passing num_classes = 1000 either, because .convnext() loads the state dict with strict = FALSE, which is silently swallowed by ...: load_state_dict() of torch for R has no strict argument.

The roxygen description had the two datasets the wrong way round ("pretrained on Imagenet 1k and fine-tuned on Imagenet 22k") and the variants table listed 21841 classes and a 252 MB file size; both are corrected (the file is 192 MB).

`convnext_small_22k1k` is pretrained on Imagenet 22k and *fine-tuned* on
Imagenet 1k, so the published weights have a 1000-class head. The formal
default was `num_classes = 21841`, taken from the other `_22k` models, so
`model_convnext_small_22k1k(pretrained = TRUE)` built a 21841-class head
and then failed to load the state dict with

  The size of tensor a (21841) must match the size of tensor b (1000)
  at non-singleton dimension 0

Verified against the published checkpoint: its `head.weight` is 1000x768,
whereas `convnext_tiny_22k` really is 21841x768.

The caller could not work around this by passing `num_classes = 1000`
either, because `.convnext()` loads the state dict with `strict = FALSE`,
which is silently swallowed by `...`: `load_state_dict()` of torch for R
has no `strict` argument.

The roxygen description had the two datasets the wrong way round
("pretrained on Imagenet 1k and fine-tuned on Imagenet 22k") and the
variants table listed 21841 classes and a 252 MB file size; both are
corrected (the file is 192 MB).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

@cregouby cregouby left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot @sebffischer !

@cregouby
cregouby merged commit 299b8a8 into mlverse:main Aug 1, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants