PyTorch implementation of the Hungarian loss function, a permutation invariant loss, with SciPy assignment problem solver. The "MIT No Attribution" (MIT-0) License.

.hungarian_loss.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from scipy.optimize import linear_sum_assignment
from torch.utils.data import DataLoader, TensorDataset


class SimpleBoundingBoxPredictor(nn.Module):
    def __init__(self):
        super(SimpleBoundingBoxPredictor, self).__init__()
        self.fc = nn.Linear(256, 4)

    def forward(self, x):
        return self.fc(x)


def match_targets(outputs, targets):
    cost_matrix = torch.cdist(outputs, targets, p=1)
    row_ind, col_ind = linear_sum_assignment(cost_matrix.cpu().detach().numpy())
    return row_ind, col_ind


def hungarian_loss(outputs, targets):
    row_ind, col_ind = match_targets(outputs, targets)
    matched_outputs = outputs[row_ind]
    matched_targets = targets[col_ind]
    loss = F.l1_loss(matched_outputs, matched_targets)
    return loss


def main():
    x_data = torch.randn(10, 256)
    y_data = torch.tensor(np.random.rand(10, 4) * 10, dtype=torch.float32)
    dataset = TensorDataset(x_data, y_data)
    data_loader = DataLoader(dataset, batch_size=2, shuffle=True)

    model = SimpleBoundingBoxPredictor()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(50):
        train_loss = 0

        for data_i, (features, targets) in enumerate(data_loader):
            optimizer.zero_grad(set_to_none=True)

            outputs = model(features)
            loss = hungarian_loss(outputs, targets)

            loss.backward()
            optimizer.step()

            train_loss += (loss.item() - train_loss) / (data_i + 1)  # Welford's online algorithm

        print(f"{epoch=}, {train_loss=}")


if __name__ == "__main__":
    main()

README.md

PyTorch Hungarian Loss Function

This repository contains a PyTorch implementation of the Hungarian loss function, utilizing the SciPy library for solving the assignment problem. The Hungarian method, also known as the Kuhn-Munkres algorithm, is an efficient way to solve assignment problems and is particularly useful in object detection tasks where predictions need to be matched with ground truth data optimally.

Features

• Simple Bounding Box Predictor: A neural network module that predicts bounding boxes given input features.
• Hungarian Loss Function: A permutation invariant loss function that computes the match between predicted bounding boxes and target bounding boxes using the Hungarian method, minimizing the L1 loss between matched pairs.
• Efficient Matching: Utilizes SciPy's linear_sum_assignment function for optimal matching between predicted and target bounding boxes.
• Dataset and DataLoader: A simple dataset setup using PyTorch's TensorDataset and DataLoader for easy batching and shuffling.

Requirements

• Python 3.6+
• PyTorch 1.7+
• SciPy 1.5+

Installation

To use this implementation, ensure you have Python installed on your system. Then, install the required libraries using pip:

pip install torch scipy

Usage

The main components of this implementation include the SimpleBoundingBoxPredictor class, the hungarian_loss function, and a training loop example. To use these components:

1. Define your dataset: Replace the synthetic dataset (x_data and y_data) with your dataset of features and bounding boxes.

2. Initialize the model and optimizer:

model = SimpleBoundingBoxPredictor()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

3. Train your model: Run the provided training loop, which automatically computes the Hungarian loss and updates the model parameters.

Example Output

During training, the script prints the epoch number and the current training loss, calculated using Welford's algorithm for a running average.

epoch=0, train_loss=2.345
epoch=1, train_loss=1.987
...

License

This project is made available under the "MIT No Attribution" (MIT-0) License. You are free to use, modify, and distribute this software without attribution.

Disclaimer

This implementation is provided as-is, without warranty. The author(s) are not responsible for any damage, data loss, or unexpected behavior resulting from the use of this software.

Feel free to contribute to the project or report any issues you encounter. Your feedback is highly appreciated!

MIT-0

MIT No Attribution

Copyright 2024 Ivan Stepanov

Permission is hereby granted, free of charge, to any person obtaining a copy of this
software and associated documentation files (the "Software"), to deal in the Software
without restriction, including without limitation the rights to use, copy, modify,
merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.