YOLOX Tiny weights don't match to the network structure

[Tomin1]

🐛 Bug

I checked the keys from state dict of YOLOX model created with the code in this project against the weights I loaded from YOLOX project for tiny sized model and I noticed that there are fewer items in the state dict.

After quite some investigation I figured that there is no corresponding layer for these (names are from YOLOX project):

backbone.backbone.dark3.1.m.2.conv1.conv.weight
backbone.backbone.dark3.1.m.2.conv1.bn.weight
backbone.backbone.dark3.1.m.2.conv1.bn.bias
backbone.backbone.dark3.1.m.2.conv1.bn.running_mean
backbone.backbone.dark3.1.m.2.conv1.bn.running_var
backbone.backbone.dark3.1.m.2.conv1.bn.num_batches_tracked
backbone.backbone.dark3.1.m.2.conv2.conv.weight
backbone.backbone.dark3.1.m.2.conv2.bn.weight
backbone.backbone.dark3.1.m.2.conv2.bn.bias
backbone.backbone.dark3.1.m.2.conv2.bn.running_mean
backbone.backbone.dark3.1.m.2.conv2.bn.running_var
backbone.backbone.dark3.1.m.2.conv2.bn.num_batches_tracked

I pinpointed this to a specific value in the code which I think is wrong when constructing YOLOX models.

To Reproduce

1. Download tiny weights from YOLOX and load it to pytorch.
2. Create YOLOXNetwork with 80 classes, width of 24 and depth of 1 which should correspond to tiny sized YOLOX network.
3. Compare the model weights (tensors) in state dictionaries (especially the number of them).

See also the code below.

Code sample

>>> import torch
>>> tiny = torch.load('yolox_tiny.pth', map_location=torch.device('cpu'), weights_only=True)["model"]
>>> len(tiny)
462
>>> import pl_bolts.models.detection
>>> yolox = pl_bolts.models.detection.yolo.torch_networks.YOLOXNetwork(80, width=24, depth=1)
>>> len(yolox.state_dict())
450
>>>

Expected behaviour

I would expect that the networks have matching number of tensors.

Environment

• PyTorch Version (e.g., 1.0): 2.6.0+cpu
• OS (e.g., Linux): Fedora Workstation 41
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source): -
• Python version: 3.13
• CUDA/cuDNN version: None
• GPU models and configuration: None
• Any other relevant information: -

Additional context

The code I figured would be the reason for this is here. If I change that depth * 2 to depth * 3, I get matching number of values (462) in the dictionary. I don't know if the current value is the same as in YOLOv5 but at least for YOLOX it is wrong apparently. I also have not checked any other sized networks, that the shapes of the tensors match, or that I could load the weights after converting the names of the layers in the state dictionary.

[Tomin1]

There might be a bit deeper issues with using that backbone with YOLOX. Tensor shapes aren't quite matching I think. At least backbone.stages.0.conv.weight shape is [24, 3, 6, 6] and the loaded backbone.backbone.stem.conv.conv.weight shape is [24, 12, 3, 3].

[Tomin1]

If I change also this line to this

downsample(width // 2, width, kernel_size=3),

then the tensor shapes match. So only a little bit tweaking was needed after all.

However, the code is not correct any more as the input size doesn't match. 😁 There would need to be some adjustments to that.