Merge pull request #189 from kozistr/feature/loss-functions

[Feature] Implement loss functions

Author: kozistr

.github/pull_request_template.md

@@ -3,11 +3,12 @@ Remove this part when you open the PR
 
 Here's a checklist before opening the Pull Request!
 
-1. PR Title convention : [Type of PR] [Summary] (e.g. [Feature] Implement AdamP optimizer)
+1. PR title convention : [Type of PR] [Summary] (e.g. [Feature] Implement AdamP optimizer)
 2. Attach `as much information as possible you can`. It helps the reviewers a lot :)
 3. Make sure the code is perfectly `runnable & compatible`.
 4. If your PR is not ready yet, make your `PR` to `Draft PR`.
-5. Make sure `make check` before opening the `PR`.
+5. Make sure `make format & check` before opening the `PR`.
+6. Or you just call the maintainer to help to fix code-style & test cases.
 ---
 
 ## Problem (Why?)

README.rst

@@ -16,7 +16,7 @@ pytorch-optimizer
 
 | **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch.
 | I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.
-| Currently, 57 optimizers, 6 lr schedulers are supported!
+| Currently, 57 optimizers, 6 lr schedulers, and 10 loss functions are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -240,6 +240,33 @@ You can check the supported learning rate schedulers with below code.
 | Chebyshev        | *Acceleration via Fractal Learning Rate Schedules*                                                |                                                                                   | `https://arxiv.org/abs/2103.01338 <https://arxiv.org/abs/2103.01338>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210301338A/exportcitation>`__                                      |
 +------------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 
+Supported Loss Function
+-----------------------
+
+You can check the supported loss functions with below code.
+
+::
+
+    from pytorch_optimizer import get_supported_loss_functions
+
+    supported_loss_functions = get_supported_loss_functions()
+
++---------------------+-------------------------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Loss Functions      | Description                                                                                                             | Official Code                                                                     | Paper                                                                                         |                                                              Citation                                                |
++=====================+=========================================================================================================================+===================================================================================+===============================================================================================+======================================================================================================================+
+| Label Smoothing     | *Rethinking the Inception Architecture for Computer Vision*                                                             |                                                                                   | `https://arxiv.org/abs/1512.00567 <https://arxiv.org/abs/1512.00567>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2015arXiv151200567S/exportcitation>`__                                      |
++---------------------+-------------------------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Focal               | *Focal Loss for Dense Object Detection*                                                                                 |                                                                                   | `https://arxiv.org/abs/1708.02002 <https://arxiv.org/abs/1708.02002>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2017arXiv170802002L/exportcitation>`__                                      |
++---------------------+-------------------------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Focal Cosine        | *Data-Efficient Deep Learning Method for Image Classification Using Data Augmentation, Focal Cosine Loss, and Ensemble* |                                                                                   | `https://arxiv.org/abs/2007.07805 <https://arxiv.org/abs/2007.07805>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv200707805K/exportcitation>`__                                      |
++---------------------+-------------------------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| LDAM                | *Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss*                                                | `github <https://github.com/kaidic/LDAM-DRW>`__                                   | `https://arxiv.org/abs/1906.07413 <https://arxiv.org/abs/1906.07413>`__                       | `cite <https://github.com/kaidic/LDAM-DRW#reference>`__                                                              |
++---------------------+-------------------------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Jaccard (IOU)       | *IoU Loss for 2D/3D Object Detection*                                                                                   |                                                                                   | `https://arxiv.org/abs/1908.03851 <https://arxiv.org/abs/1908.03851>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2019arXiv190803851Z/exportcitation>`__                                      |
++---------------------+-------------------------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Bi-Tempered         | *The Principle of Unchanged Optimality in Reinforcement Learning Generalization*                                        |                                                                                   | `https://arxiv.org/abs/1906.03361 <https://arxiv.org/abs/1906.03361>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2019arXiv190600336I/exportcitation>`__                                      |
++---------------------+-------------------------------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+
 Useful Resources
 ----------------
 

docs/changelogs/v2.11.0.md

@@ -6,6 +6,15 @@
   * [Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks](https://arxiv.org/abs/1806.06763) 
 * Implement LOMO optimizer (#188)
   * [Full Parameter Fine-tuning for Large Language Models with Limited Resources](https://arxiv.org/abs/2306.09782) 
+* Implement loss functions (#189)
+  * BCELoss
+  * BCEFocalLoss
+  * FocalLoss : [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002)
+  * FocalCosineLoss : [Data-Efficient Deep Learning Method for Image Classification Using Data Augmentation, Focal Cosine Loss, and Ensemble](https://arxiv.org/abs/2007.07805)
+  * DiceLoss : [Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations](https://arxiv.org/abs/1707.03237v3)
+  * LDAMLoss : [Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss](https://arxiv.org/abs/1906.07413)
+  * JaccardLoss
+  * BiTemperedLogisticLoss : [Robust Bi-Tempered Logistic Loss Based on Bregman Divergences](https://arxiv.org/abs/1906.03361)
 
 ### Diff
 

docs/index.rst

@@ -21,6 +21,7 @@ Contents
    base_api
    optimizer_api
    scheduler_api
+   loss_api
    util_api
 
 Indices and tables

docs/loss_api.rst

@@ -0,0 +1,82 @@
+Loss Functions
+==============
+
+.. _BCELoss:
+
+BCELoss
+-------
+
+.. autoclass:: pytorch_optimizer.BCELoss
+    :members:
+
+.. _BCEFocal:
+
+BCEFocal
+--------
+
+.. autoclass:: pytorch_optimizer.BCEFocal
+    :members:
+
+.. _FocalLoss:
+
+FocalLoss
+---------
+
+.. autoclass:: pytorch_optimizer.FocalLoss
+    :members:
+
+.. _FocalCosineLoss:
+
+FocalCosineLoss
+---------------
+
+.. autoclass:: pytorch_optimizer.FocalCosineLoss
+    :members:
+
+.. _SoftF1Loss:
+
+SoftF1Loss
+----------
+
+.. autoclass:: pytorch_optimizer.SoftF1Loss
+    :members:
+
+.. _DiceLoss:
+
+DiceLoss
+--------
+
+.. autoclass:: pytorch_optimizer.DiceLoss
+    :members:
+
+.. _LDAMLoss:
+
+LDAMLoss
+--------
+
+.. autoclass:: pytorch_optimizer.LDAMLoss
+    :members:
+
+.. _JaccardLoss:
+
+JaccardLoss
+-----------
+
+.. autoclass:: pytorch_optimizer.JaccardLoss
+    :members:
+
+.. _BiTemperedLogisticLoss:
+
+BiTemperedLogisticLoss
+----------------------
+
+.. autoclass:: pytorch_optimizer.BiTemperedLogisticLoss
+    :members:
+
+.. _BinaryBiTemperedLogisticLoss:
+
+BinaryBiTemperedLogisticLoss
+----------------------------
+
+.. autoclass:: pytorch_optimizer.BinaryBiTemperedLogisticLoss
+    :members:

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.10.1"
-description = "optimizer & lr scheduler collections in PyTorch"
+version = "2.11.0"
+description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <[email protected]>"]
 maintainers = ["kozistr <[email protected]>"]
@@ -16,7 +16,7 @@ keywords = [
     "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "GSAM", "LARS", "Lamb", "Lion", "LOMO", "Lookahead", "MADGRAD",
     "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger",
     "Ranger21", "RotoGrad", "SAM", "SGDP", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM", "SWATS", "ScalableShampoo",
-    "Shampoo", "Yogi",
+    "Shampoo", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
@@ -94,6 +94,8 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 target-version = "py311"
 
 [tool.ruff.per-file-ignores]
+"./pytorch_optimizer/__init__.py" = ["F401"]
+"./pytorch_optimizer/lr_scheduler/__init__.py" = ["F401"]
 "./hubconf.py" = ["D", "INP001"]
 "./tests/__init__.py" = ["D"]
 "./tests/constants.py" = ["D"]
@@ -104,13 +106,11 @@ target-version = "py311"
 "./tests/test_optimizers.py" = ["D", "S101"]
 "./tests/test_optimizer_parameters.py" = ["D", "S101"]
 "./tests/test_general_optimizer_parameters.py" = ["D", "S101"]
-"./tests/test_load_optimizers.py" = ["D", "S101"]
-"./tests/test_load_lr_schedulers.py" = ["D", "S101"]
 "./tests/test_lr_schedulers.py" = ["D", "S101"]
 "./tests/test_lr_scheduler_parameters.py" = ["D", "S101"]
 "./tests/test_create_optimizer.py" = ["D"]
-"./pytorch_optimizer/__init__.py" = ["F401"]
-"./pytorch_optimizer/lr_scheduler/__init__.py" = ["F401"]
+"./tests/test_loss_functions.py" = ["D", "S101"]
+"./tests/test_load_modules.py" = ["D", "S101"]
 
 [tool.pytest.ini_options]
 testpaths = "tests"

pytorch_optimizer/__init__.py

@@ -4,6 +4,13 @@
 from torch import nn
 
 from pytorch_optimizer.base.types import OPTIMIZER, PARAMETERS, SCHEDULER
+from pytorch_optimizer.loss.bi_tempered import BinaryBiTemperedLogisticLoss, BiTemperedLogisticLoss
+from pytorch_optimizer.loss.cross_entropy import BCELoss
+from pytorch_optimizer.loss.dice import DiceLoss, soft_dice_score
+from pytorch_optimizer.loss.f1 import SoftF1Loss
+from pytorch_optimizer.loss.focal import BCEFocalLoss, FocalCosineLoss, FocalLoss
+from pytorch_optimizer.loss.jaccard import JaccardLoss, soft_jaccard_score
+from pytorch_optimizer.loss.ldam import LDAMLoss
 from pytorch_optimizer.lr_scheduler import (
     ConstantLR,
     CosineAnnealingLR,
@@ -177,6 +184,22 @@
     str(lr_scheduler.__name__).lower(): lr_scheduler for lr_scheduler in LR_SCHEDULER_LIST
 }
 
+LOSS_FUNCTION_LIST: List = [
+    BCELoss,
+    BCEFocalLoss,
+    FocalLoss,
+    SoftF1Loss,
+    DiceLoss,
+    LDAMLoss,
+    FocalCosineLoss,
+    JaccardLoss,
+    BiTemperedLogisticLoss,
+    BinaryBiTemperedLogisticLoss,
+]
+LOSS_FUNCTIONS: Dict[str, nn.Module] = {
+    str(loss_function.__name__).lower(): loss_function for loss_function in LOSS_FUNCTION_LIST
+}
+
 
 def load_optimizer(optimizer: str) -> OPTIMIZER:
     optimizer: str = optimizer.lower()
@@ -245,3 +268,7 @@ def get_supported_optimizers() -> List[OPTIMIZER]:
 
 def get_supported_lr_schedulers() -> List[SCHEDULER]:
     return LR_SCHEDULER_LIST
+
+
+def get_supported_loss_functions() -> List[nn.Module]:
+    return LOSS_FUNCTION_LIST

pytorch_optimizer/base/types.py

@@ -14,3 +14,4 @@
 SCHEDULER = Type[_LRScheduler]
 
 HUTCHINSON_G = Literal['gaussian', 'rademacher']
+CLASS_MODE = Literal['binary', 'multiclass', 'multilabel']