Fix reshape_classifier_output function to correctly reshape the final output layer

[gokamisama]

Description:

This Pull Request addresses an issue in the reshape_classifier_output function, which incorrectly modifies the first nn.Linear or nn.Conv2d layer encountered in an nn.Sequential module instead of the final output layer. The proposed change ensures that the final output layer is correctly identified and reshaped to match the specified number of classes n.

Issue:

In the original implementation, when the function encounters an nn.Sequential module, it modifies the first nn.Linear or nn.Conv2d layer found. This is problematic because this layer is not necessarily the final output layer of the model, which is the one that should be reshaped.

When using VGG13 to replace the YOLOv5 model for MNIST classification, my configuration was as follows:

# yolov5/classify/train.py
...
parser.add_argument("--model", type=str, default="vgg13", help="initial weights path")
parser.add_argument("--data", type=str, default="mnist", help="cifar10, cifar100, mnist, imagenet, ...")
...

Error message:

# exception info
RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x10 and 4096x4096)

With the above configuration, the reshape_classifier_output function incorrectly modifies the first linear layer instead of the final output layer, resulting in a shape mismatch error. For example, after reshaping the classification count, the structure of VGG13 is as follows:

Solution:

The updated implementation identifies and modifies the last nn.Linear or nn.Conv2d layer within an nn.Sequential module, ensuring that the correct layer is reshaped to match the class count n.

Changes Made:

1. Updated the reshape_classifier_output function to modify the last nn.Linear or nn.Conv2d layer within an nn.Sequential module.
2. Added checks to identify the last occurrence of these layers.

Code:

def reshape_classifier_output(model, n=1000):
    """Reshapes last layer of model to match class count 'n', supporting Classify, Linear, Sequential types."""
    from models.common import Classify
    import torch.nn as nn

    name, m = list((model.model if hasattr(model, "model") else model).named_children())[-1]  # last module
    if isinstance(m, Classify):  # YOLOv5 Classify() head
        if m.linear.out_features != n:
            m.linear = nn.Linear(m.linear.in_features, n)
    elif isinstance(m, nn.Linear):  # ResNet, EfficientNet
        if m.out_features != n:
            setattr(model, name, nn.Linear(m.in_features, n))
    elif isinstance(m, nn.Sequential):
        types = [type(x) for x in m]
        if nn.Linear in types:
            # Check for Linear layers and modify the last one
            i = len(types) - 1 - types[::-1].index(nn.Linear)  # Last nn.Linear index
            if m[i].out_features != n:
                m[i] = nn.Linear(m[i].in_features, n)
        elif nn.Conv2d in types:
            # Check for Conv2d layers and modify the last one
            i = len(types) - 1 - types[::-1].index(nn.Conv2d)  # Last nn.Conv2d index
            if m[i].out_channels != n:
                m[i] = nn.Conv2d(m[i].in_channels, n, m[i].kernel_size, m[i].stride, bias=m[i].bias is not None)

Testing:

• Verified the function with different model architectures, including those with multiple nn.Linear and nn.Conv2d layers within nn.Sequential modules.
• Ensured that the final output layer is correctly reshaped to match the class count n.

Impact:

This change ensures that models using the reshape_classifier_output function correctly modify the final output layer, supporting accurate classification outputs and preventing potential shape mismatch issues.

Related Issues:

Please link any related issues or discussions here.

Acknowledgements:

Thank you to the maintainers for reviewing this Pull Request. I greatly appreciate your feedback.

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Improvements in neural network classifier output reshaping.

📊 Key Changes

• The method to identify the index of nn.Linear and nn.Conv2d layers in a network has been updated to find the last occurrence of these layers instead of the first.

🎯 Purpose & Impact

• Purpose: To ensure that when the output layer of a neural network (either nn.Linear or nn.Conv2d) is being reshaped to match a specified number of outputs (n=1000), the last layer is modified. This is crucial for networks where these layers might appear multiple times.
• Impact: This change allows for more accurate customizations of neural network architectures, particularly in cases where the final layer needs to be adjusted for different numbers of output classes. It enhances model flexibility and adaptability for a variety of tasks.
• Users: This adjustment is essential for developers and researchers working on custom object detection models, enabling them to tailor their models more precisely to specific needs. Non-expert users benefit indirectly through improved model accuracy and performance in applications powered by these customized models. 🚀

[github-actions]

All Contributors have signed the CLA. ✅
_{Posted by the CLA Assistant Lite bot.}

[github-actions]

👋 Hello @goksmisama, thank you for submitting a YOLOv5 🚀 PR! To allow your work to be integrated as seamlessly as possible, we advise you to:

• ✅ Verify your PR is up-to-date with ultralytics/yolov5 master branch. If your PR is behind you can update your code by clicking the 'Update branch' button or by running git pull and git merge master locally.
• ✅ Verify all YOLOv5 Continuous Integration (CI) checks are passing.
• ✅ Reduce changes to the absolute minimum required for your bug fix or feature addition. "It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is." — Bruce Lee

[gokamisama]

I have read the CLA Document and I sign the CLA

[gokamisama]

@glenn-jocher Could you please review this PR when you have time? Thank you!

[glenn-jocher]

Hi @goksmisama, just a gentle reminder to review this PR at your earliest convenience. Thanks for your time! 😊

[gokamisama]

Hi @glenn-jocher, thanks for your comment! I have reviewed the PR. Please let me know if there is anything else I need to do. Thanks again!

[glenn-jocher]

@goksmisama PR merged! Thank you for your contributions :)