chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,382 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=redefined-builtin, invalid-name
|
||||
# ruff: noqa: RUF012
|
||||
"""Loss functions library for relax."""
|
||||
|
||||
# isort: off
|
||||
from typing import Literal
|
||||
|
||||
# isort: on
|
||||
|
||||
from ..block_builder import BlockBuilder
|
||||
from ..expr import Expr, Function, Type, Var
|
||||
from ..op import abs, argmax, mean, multiply, reshape, subtract, sum
|
||||
from ..op.nn import log_softmax, nll_loss
|
||||
|
||||
|
||||
def _create_param_var(param: Var | Type, param_name: str) -> Var:
|
||||
"""If param is a Type, create a Var with the given Type and name.
|
||||
|
||||
If param is a Var, create a Var with the same Type and name as the given param Var."""
|
||||
if isinstance(param, Type):
|
||||
param = Var(param_name, param)
|
||||
if not isinstance(param, Var):
|
||||
raise TypeError("The type of param should be Var or Type, but got " + type(param))
|
||||
return Var(param.name_hint, param.ty)
|
||||
|
||||
|
||||
class Loss:
|
||||
r"""Base class of all loss.
|
||||
|
||||
Generally, loss function will take one or more **input parameters** (that is outputs of
|
||||
the backbone of a model), one or more **target parameters**, and generate a scalar value
|
||||
denoting the loss.
|
||||
|
||||
You can use `relax.transform.AppendLoss` to append the loss function to a one-dataflowblock
|
||||
backbone function in a IRModule. That will generate a one-dataflowblock function accepting
|
||||
instances and targets, and then returning the loss.
|
||||
|
||||
Most loss functions involve a reduction of losses from all instances in a batch. We use
|
||||
`reduction` parameter to denote the reduction method. Possible reduction methods include
|
||||
`"mean"`, `"sum"` and `"none"`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
loss_name : str
|
||||
The name of the loss function. Should be provided when calling `super().__init__` in
|
||||
constructor functions of subclasses.
|
||||
|
||||
num_backbone_outputs : int
|
||||
The number of `prediction_outputs` of the backbone function, alos the number of the
|
||||
backbone_prediction_outputs of the loss function. See `relax.transform.AppendLoss`.
|
||||
|
||||
Should be provided when calling `super().__init__` in constructor functions of subclasses.
|
||||
|
||||
For example, `CrossEntropyLoss` requires one backbone prediction output; `MarginRankingLoss`
|
||||
requires two backbone prediction outputs.
|
||||
|
||||
reduction : Literal["mean", "sum", "none"]
|
||||
The reduction method to apply to output. Can be "mean", "sum" or "none".
|
||||
|
||||
none : no reduction will be applied,
|
||||
mean : the sum of the output will be divided by the batch_size,
|
||||
sum : the output will be summed.
|
||||
"""
|
||||
|
||||
_valid_reductions = ["mean", "sum", "none"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
loss_name: str,
|
||||
num_backbone_outputs: int,
|
||||
reduction: Literal["mean", "sum", "none"] = "mean",
|
||||
) -> None:
|
||||
self._loss_name = loss_name
|
||||
self._reduction = reduction
|
||||
self._num_backbone_outputs = num_backbone_outputs
|
||||
|
||||
if self._reduction not in self._valid_reductions:
|
||||
raise ValueError("Reduction can only be one of these values: ", self._valid_reductions)
|
||||
|
||||
@property
|
||||
def num_backbone_outputs(self) -> int:
|
||||
"""Get the number of number of the outputs of the backbone function."""
|
||||
return self._num_backbone_outputs
|
||||
|
||||
def _with_reduction(self, expr: Expr) -> Expr:
|
||||
"""Add a reduction to the final loss.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
expr : Expr
|
||||
The loss expr.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : Expr
|
||||
The reduced result.
|
||||
"""
|
||||
if self._reduction == "sum":
|
||||
expr = sum(expr)
|
||||
elif self._reduction == "mean":
|
||||
expr = mean(expr)
|
||||
elif self._reduction != "none":
|
||||
raise ValueError("Reduction can only be one of these values: ", self._valid_reductions)
|
||||
return expr
|
||||
|
||||
|
||||
class L1Loss(Loss):
|
||||
r"""Mean element-wise absolute value difference.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
reduction : Literal["mean", "sum", "none"]
|
||||
The reduction method to apply to output. Can be "mean", "sum" or "none".
|
||||
|
||||
none : no reduction will be applied,
|
||||
mean : the sum of the output will be divided by the batch_size,
|
||||
sum : the output will be summed.
|
||||
"""
|
||||
|
||||
def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean") -> None:
|
||||
super().__init__("l1_loss", 1, reduction)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
predictions: Var | Type,
|
||||
targets: Var | Type,
|
||||
) -> Function:
|
||||
"""Get the relax function of L1Loss. If the parameters are
|
||||
type, it will create corresponding variables.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
predictions : Union[Var, Type]
|
||||
The predictions of the model in the calculation of loss.
|
||||
targets : Union[Var, Type]
|
||||
The ground truth in the calculation of loss.
|
||||
|
||||
Returns
|
||||
-------
|
||||
The relax function of L1Loss with the loss name as its global symbol.
|
||||
"""
|
||||
bb = BlockBuilder()
|
||||
|
||||
predictions = _create_param_var(predictions, "predictions")
|
||||
targets = _create_param_var(targets, "targets")
|
||||
|
||||
with bb.function(self._loss_name, [predictions, targets]):
|
||||
with bb.dataflow():
|
||||
lv = abs(subtract(predictions, targets))
|
||||
loss = bb.emit_output(self._with_reduction(lv))
|
||||
bb.emit_func_output(loss)
|
||||
|
||||
return bb.get()[self._loss_name]
|
||||
|
||||
|
||||
class MSELoss(Loss):
|
||||
r"""Measures the element-wise mean squared error.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
reduction : Literal["mean", "sum", "none"]
|
||||
The reduction method to apply to output. Can be "mean", "sum" or "none".
|
||||
|
||||
none : no reduction will be applied,
|
||||
mean : the sum of the output will be divided by the batch_size,
|
||||
sum : the output will be summed.
|
||||
"""
|
||||
|
||||
def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean") -> None:
|
||||
super().__init__("mse_loss", 1, reduction)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
predictions: Var | Type,
|
||||
targets: Var | Type,
|
||||
) -> Function:
|
||||
"""Get the relax function of MSELoss. If the parameters are
|
||||
type, it will create corresponding variables.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
predictions : Union[Var, Type]
|
||||
The predictions of the model in the calculation of loss.
|
||||
targets : Union[Var, Type]
|
||||
The ground truth in the calculation of loss.
|
||||
|
||||
Returns
|
||||
-------
|
||||
The relax function of MSELoss with the loss name as its global symbol.
|
||||
"""
|
||||
bb = BlockBuilder()
|
||||
|
||||
predictions = _create_param_var(predictions, "predictions")
|
||||
targets = _create_param_var(targets, "targets")
|
||||
|
||||
with bb.function(self._loss_name, [predictions, targets]):
|
||||
with bb.dataflow():
|
||||
lv = subtract(predictions, targets)
|
||||
lv = multiply(lv, lv)
|
||||
loss = bb.emit_output(self._with_reduction(lv))
|
||||
bb.emit_func_output(loss)
|
||||
|
||||
return bb.get()[self._loss_name]
|
||||
|
||||
|
||||
class CrossEntropyLoss(Loss):
|
||||
r"""CrossEntropyLoss. It is a combination of a log_softmax computation and a nll_loss.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
reduction : Literal["mean", "sum", "none"]
|
||||
The reduction method to apply to output. Can be "mean", "sum" or "none".
|
||||
|
||||
none : no reduction will be applied,
|
||||
mean : the sum of the output will be divided by the batch_size,
|
||||
sum : the output will be summed.
|
||||
|
||||
ignore_index : int
|
||||
Specifies a target value that is ignored and does not contribute to the input gradient.
|
||||
"""
|
||||
|
||||
ignore_index: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reduction: Literal["mean", "sum", "none"] = "mean",
|
||||
ignore_index: int = -100,
|
||||
) -> None:
|
||||
super().__init__("cross_entropy_loss", 1, reduction)
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
predictions: Var | Type,
|
||||
targets: Var | Type,
|
||||
weights: Var | Type | None = None,
|
||||
) -> Function:
|
||||
"""Get the relax function of CrossEntropyLoss. If the parameters are
|
||||
type, it will create corresponding variables.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
predictions : Union[Var, Type]
|
||||
The predictions of the model in the calculation of loss.
|
||||
|
||||
targets : Union[Var, Type]
|
||||
The ground truth in the calculation of loss.
|
||||
|
||||
weights : Optional[Union[Var, Type]]
|
||||
a manual rescaling weight given to each class. It has to be a Tensor of size C.
|
||||
|
||||
Returns
|
||||
-------
|
||||
The relax function of CrossEntropyLoss with the loss name as its global symbol.
|
||||
"""
|
||||
bb = BlockBuilder()
|
||||
|
||||
predictions = _create_param_var(predictions, "predictions")
|
||||
targets = _create_param_var(targets, "targets")
|
||||
|
||||
arg_list = [predictions, targets]
|
||||
if weights:
|
||||
weights = _create_param_var(weights, "weights")
|
||||
arg_list.append(weights)
|
||||
|
||||
with bb.function(self._loss_name, arg_list):
|
||||
with bb.dataflow():
|
||||
logits = bb.emit(log_softmax(predictions))
|
||||
loss = bb.emit_output(
|
||||
nll_loss(logits, targets, weights, self._reduction, self.ignore_index)
|
||||
)
|
||||
bb.emit_func_output(loss)
|
||||
|
||||
return bb.get()[self._loss_name]
|
||||
|
||||
|
||||
class CategoricalCrossEntropyLoss(Loss):
|
||||
r"""CategoricalCrossEntropyLoss.
|
||||
It is a combination of a converting one-hot target vector to a label,
|
||||
a log_softmax computation and a nll_loss.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
reduction : Literal["mean", "sum", "none"]
|
||||
The reduction method to apply to output. Can be "mean", "sum" or "none".
|
||||
|
||||
none : no reduction will be applied,
|
||||
mean : the sum of the output will be divided by the batch_size,
|
||||
sum : the output will be summed.
|
||||
|
||||
ignore_index : int
|
||||
Specifies a target value that is ignored and does not contribute to the input gradient.
|
||||
"""
|
||||
|
||||
ignore_index: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reduction: Literal["mean", "sum", "none"] = "mean",
|
||||
ignore_index: int = -100,
|
||||
) -> None:
|
||||
super().__init__("categorical_cross_entropy_loss", 1, reduction)
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
predictions: Var | Type,
|
||||
targets: Var | Type,
|
||||
weights: Var | Type | None = None,
|
||||
) -> Function:
|
||||
"""Get the relax function of CategoricalCrossEntropyLoss. If the parameters are
|
||||
type, it will create corresponding variables.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
predictions : Union[Var, Type]
|
||||
The predictions of the model in the calculation of loss.
|
||||
|
||||
targets : Union[Var, Type]
|
||||
The ground truth in the calculation of loss.
|
||||
|
||||
weights : Optional[Union[Var, Type]]
|
||||
a manual rescaling weight given to each class. It has to be a Tensor of size C.
|
||||
|
||||
Returns
|
||||
-------
|
||||
The relax function of CategoricalCrossEntropyLoss with the loss name as its global symbol.
|
||||
"""
|
||||
|
||||
if "int" not in str(targets.dtype):
|
||||
raise TypeError(
|
||||
f"Dtype of targets expected to be int/uint. \
|
||||
However, the dtype of targets is {targets.dtype}"
|
||||
)
|
||||
|
||||
bb = BlockBuilder()
|
||||
|
||||
predictions = _create_param_var(predictions, "predictions")
|
||||
targets = _create_param_var(targets, "targets")
|
||||
|
||||
arg_list = [predictions, targets]
|
||||
if weights:
|
||||
weights = _create_param_var(weights, "weights")
|
||||
arg_list.append(weights)
|
||||
|
||||
# In the case of ignore_index >= 0,
|
||||
# the nll_loss function is used to handle the ignore index.
|
||||
# In other cases where ignore_index is not needed, just use the simpe product.
|
||||
with bb.function(self._loss_name, arg_list):
|
||||
with bb.dataflow():
|
||||
logits = bb.emit(log_softmax(predictions))
|
||||
if self.ignore_index >= 0:
|
||||
targets = bb.emit(
|
||||
reshape(argmax(targets, axis=1), shape=(targets.ty.shape[0],))
|
||||
)
|
||||
loss = bb.emit_output(
|
||||
nll_loss(logits, targets, weights, self._reduction, self.ignore_index)
|
||||
)
|
||||
else:
|
||||
lv = bb.emit(-logits * targets.astype("float32"))
|
||||
if weights:
|
||||
lv = bb.emit(lv * weights)
|
||||
loss = bb.emit_output(self._with_reduction(lv))
|
||||
bb.emit_func_output(loss)
|
||||
|
||||
return bb.get()[self._loss_name]
|
||||
Reference in New Issue
Block a user