损失函数

目录

损失函数#

import tvm.testing
from tvm import relax
from tvm.ir.base import assert_structural_equal
from tvm.script import relax as R, ir as I

@I.ir_module
class Module:
    @R.function
    def forward(
        x: R.Tensor((2, 4), "float32"),
        w: R.Tensor((4, 4), "float32"),
        b: R.Tensor((2, 4), "float32"),
    ) -> R.Tensor((2, 4), "float32"):
        with R.dataflow():
            lv: R.Tensor((2, 4), "float32") = R.matmul(x, w)
            out: R.Tensor((2, 4), "float32") = R.add(lv, b)
            R.output(out)
        return out

L1 损失#

N = 3
C = 5
predictions = relax.TensorStructInfo((N, C), "float32")
targets = relax.TensorStructInfo((N, C), "float32")
l1_loss = relax.training.loss.L1Loss()
l1_loss(predictions, targets).show()

# from tvm.script import relax as R

@R.function
def l1_loss(predictions: R.Tensor((3, 5), dtype="float32"), targets: R.Tensor((3, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((3, 5), dtype="float32") = R.subtract(predictions, targets)
        lv1: R.Tensor((3, 5), dtype="float32") = R.abs(lv)
        gv: R.Tensor((), dtype="float32") = R.mean(lv1, axis=None, keepdims=False)
        R.output(gv)
    return gv

s = Module["forward"].ret_struct_info
l1_loss = relax.training.loss.L1Loss(reduction="sum")
After = relax.training.AppendLoss("forward", l1_loss(s, s), l1_loss.num_backbone_outputs)(
    Module
)
After["forward_loss"].show()

# from tvm.script import relax as R

@R.function
def forward_loss(x: R.Tensor((2, 4), dtype="float32"), w: R.Tensor((4, 4), dtype="float32"), b: R.Tensor((2, 4), dtype="float32"), targets: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((2, 4), dtype="float32") = R.matmul(x, w, out_dtype="void")
        out: R.Tensor((2, 4), dtype="float32") = R.add(lv, b)
        lv_1: R.Tensor((2, 4), dtype="float32") = R.subtract(out, targets)
        lv1: R.Tensor((2, 4), dtype="float32") = R.abs(lv_1)
        gv: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
        R.output(gv)
    return gv

MSE 损失#

N = 3
C = 5
predictions = relax.TensorStructInfo((N, C), "float32")
targets = relax.TensorStructInfo((N, C), "float32")
mse_loss = relax.training.loss.MSELoss()
mse_loss(predictions, targets).show()

# from tvm.script import relax as R

@R.function
def mse_loss(predictions: R.Tensor((3, 5), dtype="float32"), targets: R.Tensor((3, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((3, 5), dtype="float32") = R.subtract(predictions, targets)
        lv1: R.Tensor((3, 5), dtype="float32") = R.multiply(lv, lv)
        gv: R.Tensor((), dtype="float32") = R.mean(lv1, axis=None, keepdims=False)
        R.output(gv)
    return gv

s = Module["forward"].ret_struct_info
mse_loss = relax.training.loss.MSELoss(reduction="sum")
After = relax.training.AppendLoss("forward", mse_loss(s, s), mse_loss.num_backbone_outputs)(
    Module
)
After["forward_loss"].show()

# from tvm.script import relax as R

@R.function
def forward_loss(x: R.Tensor((2, 4), dtype="float32"), w: R.Tensor((4, 4), dtype="float32"), b: R.Tensor((2, 4), dtype="float32"), targets: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((2, 4), dtype="float32") = R.matmul(x, w, out_dtype="void")
        out: R.Tensor((2, 4), dtype="float32") = R.add(lv, b)
        lv_1: R.Tensor((2, 4), dtype="float32") = R.subtract(out, targets)
        lv1: R.Tensor((2, 4), dtype="float32") = R.multiply(lv_1, lv_1)
        gv: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
        R.output(gv)
    return gv

交叉熵损失#

带有权重：

N = 3
C = 5
predictions = relax.TensorStructInfo((N, C), "float32")
targets = relax.TensorStructInfo((N,), "int64")
weights = relax.TensorStructInfo((C,), "float32")
cross_entropy_loss = relax.training.loss.CrossEntropyLoss(reduction="sum", ignore_index=1)
cross_entropy_loss(predictions, targets, weights).show()

# from tvm.script import relax as R

@R.function
def cross_entropy_loss(predictions: R.Tensor((3, 5), dtype="float32"), targets: R.Tensor((3,), dtype="int64"), weights: R.Tensor((5,), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((3, 5), dtype="float32") = R.nn.log_softmax(predictions, axis=-1)
        gv: R.Tensor((), dtype="float32") = R.nn.nll_loss(lv, targets, weights, reduction="sum", ignore_index=1)
        R.output(gv)
    return gv

无权重：

N = 3
C = 5
predictions = relax.TensorStructInfo((N, C), "float32")
targets = relax.TensorStructInfo((N,), "int64")
cross_entropy_loss = relax.training.loss.CrossEntropyLoss()
cross_entropy_loss(predictions, targets).show()

# from tvm.script import relax as R

@R.function
def cross_entropy_loss(predictions: R.Tensor((3, 5), dtype="float32"), targets: R.Tensor((3,), dtype="int64")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((3, 5), dtype="float32") = R.nn.log_softmax(predictions, axis=-1)
        gv: R.Tensor((), dtype="float32") = R.nn.nll_loss(lv, targets, reduction="mean", ignore_index=-100)
        R.output(gv)
    return gv

s = Module["forward"].ret_struct_info
N = s.shape[0]
C = s.shape[1]
targets = relax.TensorStructInfo((N,), "int64")
weights = relax.TensorStructInfo((C,), "float32")
cross_entropy_loss = relax.training.loss.CrossEntropyLoss(reduction="sum", ignore_index=1)
After = relax.training.AppendLoss(
    "forward", cross_entropy_loss(s, targets, weights), cross_entropy_loss.num_backbone_outputs
)(Module)
After["forward_loss"].show()

# from tvm.script import relax as R

@R.function
def forward_loss(x: R.Tensor((2, 4), dtype="float32"), w: R.Tensor((4, 4), dtype="float32"), b: R.Tensor((2, 4), dtype="float32"), targets: R.Tensor((2,), dtype="int64"), weights: R.Tensor((4,), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((2, 4), dtype="float32") = R.matmul(x, w, out_dtype="void")
        out: R.Tensor((2, 4), dtype="float32") = R.add(lv, b)
        lv_1: R.Tensor((2, 4), dtype="float32") = R.nn.log_softmax(out, axis=-1)
        gv: R.Tensor((), dtype="float32") = R.nn.nll_loss(lv_1, targets, weights, reduction="sum", ignore_index=1)
        R.output(gv)
    return gv

分类交叉熵损失#

N = 3
C = 5
predictions = relax.TensorStructInfo((N, C), "float32")
targets = relax.TensorStructInfo((N, C), "int64")
weights = relax.TensorStructInfo((C,), "float32")
categorical_cross_entropy_loss = relax.training.loss.CategoricalCrossEntropyLoss(
    reduction="sum"
)
categorical_cross_entropy_loss(predictions, targets, weights).show()

# from tvm.script import relax as R

@R.function
def categorical_cross_entropy_loss(predictions: R.Tensor((3, 5), dtype="float32"), targets: R.Tensor((3, 5), dtype="int64"), weights: R.Tensor((5,), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((3, 5), dtype="float32") = R.nn.log_softmax(predictions, axis=-1)
        lv1: R.Tensor((3, 5), dtype="float32") = R.negative(lv)
        lv2: R.Tensor((3, 5), dtype="float32") = R.astype(targets, dtype="float32")
        lv3: R.Tensor((3, 5), dtype="float32") = R.multiply(lv1, lv2)
        lv4: R.Tensor((3, 5), dtype="float32") = R.multiply(lv3, weights)
        gv: R.Tensor((), dtype="float32") = R.sum(lv4, axis=None, keepdims=False)
        R.output(gv)
    return gv

N = 3
C = 5
predictions = relax.TensorStructInfo((N, C), "float32")
targets = relax.TensorStructInfo((N, C), "int64")
categorical_cross_entropy_loss = relax.training.loss.CategoricalCrossEntropyLoss()
categorical_cross_entropy_loss(predictions, targets).show()

# from tvm.script import relax as R

@R.function
def categorical_cross_entropy_loss(predictions: R.Tensor((3, 5), dtype="float32"), targets: R.Tensor((3, 5), dtype="int64")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((3, 5), dtype="float32") = R.nn.log_softmax(predictions, axis=-1)
        lv1: R.Tensor((3, 5), dtype="float32") = R.negative(lv)
        lv2: R.Tensor((3, 5), dtype="float32") = R.astype(targets, dtype="float32")
        lv3: R.Tensor((3, 5), dtype="float32") = R.multiply(lv1, lv2)
        gv: R.Tensor((), dtype="float32") = R.mean(lv3, axis=None, keepdims=False)
        R.output(gv)
    return gv

忽略目标值为 1 的类别：

N = 3
C = 5
predictions = relax.TensorStructInfo((N, C), "float32")
targets = relax.TensorStructInfo((N, C), "int64")
weights = relax.TensorStructInfo((C,), "float32")
categorical_cross_entropy_loss = relax.training.loss.CategoricalCrossEntropyLoss(
    reduction="sum", ignore_index=1
)
categorical_cross_entropy_loss(predictions, targets, weights).show()

# from tvm.script import relax as R

@R.function
def categorical_cross_entropy_loss(predictions: R.Tensor((3, 5), dtype="float32"), targets: R.Tensor((3, 5), dtype="int64"), weights: R.Tensor((5,), dtype="float32")) -> R.Tensor((), dtype="float32"):
    with R.dataflow():
        lv: R.Tensor((3, 5), dtype="float32") = R.nn.log_softmax(predictions, axis=-1)
        lv1: R.Tensor((3,), dtype="int64") = R.argmax(targets, axis=1, keepdims=False)
        lv2: R.Tensor((3,), dtype="int64") = R.reshape(lv1, R.shape([3]))
        gv: R.Tensor((), dtype="float32") = R.nn.nll_loss(lv, lv2, weights, reduction="sum", ignore_index=1)
        R.output(gv)
    return gv