chore: import upstream snapshot with attribution

tests/python/relax/test_training_loss.py

@@ -0,0 +1,299 @@
+# 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.
+import tvm.testing
+from tvm import relax
+from tvm.ir.base import assert_structural_equal
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+@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
+
+
+def test_l1_loss():
+    N = 3
+    C = 5
+    predictions = relax.TensorType((N, C), "float32")
+    targets = relax.TensorType((N, C), "float32")
+    l1_loss = relax.training.loss.L1Loss()
+
+    @R.function
+    def expected(
+        predictions: R.Tensor((3, 5), "float32"), targets: R.Tensor((3, 5), "float32")
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "l1_loss"})
+        with R.dataflow():
+            lv: R.Tensor((3, 5), "float32") = R.subtract(predictions, targets)
+            lv1: R.Tensor((3, 5), "float32") = R.abs(lv)
+            gv: R.Tensor((), "float32") = R.mean(lv1, axis=None, keepdims=False)
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(l1_loss(predictions, targets), expected)
+
+
+def test_l1_loss_append():
+    s = Module["forward"].ret_ty
+    l1_loss = relax.training.loss.L1Loss(reduction="sum")
+    After = relax.training.AppendLoss("forward", l1_loss(s, s), l1_loss.num_backbone_outputs)(
+        Module
+    )
+
+    @R.function
+    def expected(
+        x: R.Tensor((2, 4), "float32"),
+        w: R.Tensor((4, 4), "float32"),
+        b: R.Tensor((2, 4), "float32"),
+        targets: R.Tensor((2, 4), "float32"),
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "forward_loss"})
+        with R.dataflow():
+            lv: R.Tensor((2, 4), "float32") = R.matmul(x, w)
+            out: R.Tensor((2, 4), "float32") = R.add(lv, b)
+            lv1: R.Tensor((2, 4), "float32") = R.subtract(out, targets)
+            lv11: R.Tensor((2, 4), "float32") = R.abs(lv1)
+            gv: R.Tensor((), "float32") = R.sum(lv11, axis=None, keepdims=False)
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(After["forward_loss"], expected)
+
+
+def test_mse_loss():
+    N = 3
+    C = 5
+    predictions = relax.TensorType((N, C), "float32")
+    targets = relax.TensorType((N, C), "float32")
+    mse_loss = relax.training.loss.MSELoss()
+
+    @R.function
+    def expected(
+        predictions: R.Tensor((3, 5), "float32"), targets: R.Tensor((3, 5), "float32")
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "mse_loss"})
+        with R.dataflow():
+            lv: R.Tensor((3, 5), "float32") = R.subtract(predictions, targets)
+            lv1: R.Tensor((3, 5), "float32") = R.multiply(lv, lv)
+            gv: R.Tensor((), "float32") = R.mean(lv1, axis=None, keepdims=False)
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(mse_loss(predictions, targets), expected)
+
+
+def test_mse_loss_append():
+    s = Module["forward"].ret_ty
+    mse_loss = relax.training.loss.MSELoss(reduction="sum")
+    After = relax.training.AppendLoss("forward", mse_loss(s, s), mse_loss.num_backbone_outputs)(
+        Module
+    )
+
+    @R.function
+    def expected(
+        x: R.Tensor((2, 4), "float32"),
+        w: R.Tensor((4, 4), "float32"),
+        b: R.Tensor((2, 4), "float32"),
+        targets: R.Tensor((2, 4), "float32"),
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "forward_loss"})
+        with R.dataflow():
+            lv: R.Tensor((2, 4), "float32") = R.matmul(x, w)
+            out: R.Tensor((2, 4), "float32") = R.add(lv, b)
+            lv1: R.Tensor((2, 4), "float32") = R.subtract(out, targets)
+            lv11: R.Tensor((2, 4), "float32") = R.multiply(lv1, lv1)
+            gv: R.Tensor((), "float32") = R.sum(lv11, axis=None, keepdims=False)
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(After["forward_loss"], expected)
+
+
+def test_cross_entropy_loss():
+    N = 3
+    C = 5
+    predictions = relax.TensorType((N, C), "float32")
+    targets = relax.TensorType((N,), "int64")
+    weights = relax.TensorType((C,), "float32")
+    cross_entropy_loss = relax.training.loss.CrossEntropyLoss(reduction="sum", ignore_index=1)
+
+    @R.function
+    def expected(
+        predictions: R.Tensor((3, 5), "float32"),
+        targets: R.Tensor((3,), "int64"),
+        weights: R.Tensor((5,), "float32"),
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "cross_entropy_loss"})
+        with R.dataflow():
+            lv: R.Tensor((3, 5), "float32") = R.nn.log_softmax(predictions, axis=-1)
+            gv: R.Tensor((), "float32") = R.nn.nll_loss(
+                lv, targets, weights, reduction="sum", ignore_index=1
+            )
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(cross_entropy_loss(predictions, targets, weights), expected)
+
+
+def test_cross_entropy_loss_without_weights():
+    N = 3
+    C = 5
+    predictions = relax.TensorType((N, C), "float32")
+    targets = relax.TensorType((N,), "int64")
+    cross_entropy_loss = relax.training.loss.CrossEntropyLoss()
+
+    @R.function
+    def expected(
+        predictions: R.Tensor((3, 5), "float32"), targets: R.Tensor((3,), "int64")
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "cross_entropy_loss"})
+        with R.dataflow():
+            lv: R.Tensor((3, 5), "float32") = R.nn.log_softmax(predictions, axis=-1)
+            gv: R.Tensor((), "float32") = R.nn.nll_loss(
+                lv, targets, reduction="mean", ignore_index=-100
+            )
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(cross_entropy_loss(predictions, targets), expected)
+
+
+def test_cross_entropy_loss_append():
+    s = Module["forward"].ret_ty
+    N = s.shape[0]
+    C = s.shape[1]
+    targets = relax.TensorType((N,), "int64")
+    weights = relax.TensorType((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)
+
+    @R.function
+    def expected(
+        x: R.Tensor((2, 4), "float32"),
+        w: R.Tensor((4, 4), "float32"),
+        b: R.Tensor((2, 4), "float32"),
+        targets: R.Tensor((2,), "int64"),
+        weights: R.Tensor((4,), "float32"),
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "forward_loss"})
+        with R.dataflow():
+            lv: R.Tensor((2, 4), "float32") = R.matmul(x, w)
+            out: R.Tensor((2, 4), "float32") = R.add(lv, b)
+            lv1: R.Tensor((2, 4), "float32") = R.nn.log_softmax(out, axis=-1)
+            gv: R.Tensor((), "float32") = R.nn.nll_loss(
+                lv1, targets, weights, reduction="sum", ignore_index=1
+            )
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(After["forward_loss"], expected)
+
+
+def test_categorical_cross_entropy_loss():
+    N = 3
+    C = 5
+    predictions = relax.TensorType((N, C), "float32")
+    targets = relax.TensorType((N, C), "int64")
+    weights = relax.TensorType((C,), "float32")
+    categorical_cross_entropy_loss = relax.training.loss.CategoricalCrossEntropyLoss(
+        reduction="sum"
+    )
+
+    @R.function
+    def expected(
+        predictions: R.Tensor((3, 5), "float32"),
+        targets: R.Tensor((3, 5), "int64"),
+        weights: R.Tensor((5,), "float32"),
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "categorical_cross_entropy_loss"})
+        with R.dataflow():
+            lv: R.Tensor((3, 5), "float32") = R.nn.log_softmax(predictions, axis=-1)
+            lv: R.Tensor((), "float32") = -lv * targets.astype("float32")
+            gv: R.Tensor((), "float32") = R.sum(lv * weights)
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(categorical_cross_entropy_loss(predictions, targets, weights), expected)
+
+
+def test_categorical_cross_entropy_loss_without_weights():
+    N = 3
+    C = 5
+    predictions = relax.TensorType((N, C), "float32")
+    targets = relax.TensorType((N, C), "int64")
+    categorical_cross_entropy_loss = relax.training.loss.CategoricalCrossEntropyLoss()
+
+    @R.function
+    def expected(
+        predictions: R.Tensor((3, 5), "float32"), targets: R.Tensor((3, 5), "int64")
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "categorical_cross_entropy_loss"})
+        with R.dataflow():
+            lv: R.Tensor((3, 5), "float32") = R.nn.log_softmax(predictions, axis=-1)
+            gv: R.Tensor((), "float32") = R.mean(-lv * targets.astype("float32"))
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(categorical_cross_entropy_loss(predictions, targets), expected)
+
+
+def test_categorical_cross_entropy_loss_with_ignore_index():
+    N = 3
+    C = 5
+    predictions = relax.TensorType((N, C), "float32")
+    targets = relax.TensorType((N, C), "int64")
+    weights = relax.TensorType((C,), "float32")
+    categorical_cross_entropy_loss = relax.training.loss.CategoricalCrossEntropyLoss(
+        reduction="sum", ignore_index=1
+    )
+
+    @R.function
+    def expected(
+        predictions: R.Tensor((3, 5), "float32"),
+        targets: R.Tensor((3, 5), "int64"),
+        weights: R.Tensor((5,), "float32"),
+    ) -> R.Tensor((), "float32"):
+        R.func_attr({"global_symbol": "categorical_cross_entropy_loss"})
+        with R.dataflow():
+            lv: R.Tensor((3, 5), "float32") = R.nn.log_softmax(predictions, axis=-1)
+            targets = relax.op.reshape(
+                relax.op.argmax(targets, axis=1), shape=(targets.ty.shape[0],)
+            )
+            gv: R.Tensor((), "float32") = R.nn.nll_loss(
+                lv, targets, weights, reduction="sum", ignore_index=1
+            )
+            R.output(gv)
+        return gv
+
+    assert_structural_equal(categorical_cross_entropy_loss(predictions, targets, weights), expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()