# 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. # ruff: noqa: E501 import pytest import tvm import tvm.testing from tvm import relax from tvm.ir.base import assert_structural_equal from tvm.relax.training import SetupTrainer from tvm.relax.training.loss import MSELoss from tvm.relax.training.optimizer import SGD, MomentumSGD from tvm.script import ir as I from tvm.script import relax as R def test_simple(): # fmt: off @I.ir_module class Backbone: I.module_attrs({"param_num": 1, "state_num": 0}) @R.function def backbone(x: R.Tensor((2, 2), "float64"), y: R.Tensor((2, 2), "float64")): with R.dataflow(): x1 = x + y R.output(x1) return x1 @I.ir_module class Expected: I.module_attrs({"input_num": 1, "param_num": 1, "state_num": 0}) @R.function def backbone(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64")) -> R.Tensor((2, 2), dtype="float64"): with R.dataflow(): x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y) R.output(x1) return x1 @R.function def backbone_loss(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64"), targets: R.Tensor((2, 2), dtype="float64")) -> R.Tensor((), dtype="float64"): with R.dataflow(): x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y) lv: R.Tensor((2, 2), dtype="float64") = R.subtract(x1, targets) lv1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv, lv) gv: R.Tensor((), dtype="float64") = R.sum(lv1, axis=None, keepdims=False) R.output(gv) return gv @R.function def backbone_loss_adjoint(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64"), targets: R.Tensor((2, 2), dtype="float64")) -> R.Tuple(R.Tensor((), dtype="float64"), R.Tuple(R.Tensor((2, 2), dtype="float64"))): with R.dataflow(): x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y) lv: R.Tensor((2, 2), dtype="float64") = R.subtract(x1, targets) lv1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv, lv) gv: R.Tensor((), dtype="float64") = R.sum(lv1, axis=None, keepdims=False) gv_adjoint: R.Tensor((), dtype="float64") = R.ones(R.shape([]), dtype="float64") lv1_adjoint: R.Tensor((2, 2), dtype="float64") = R.broadcast_to(gv_adjoint, R.shape([2, 2])) lv_adjoint: R.Tensor((2, 2), dtype="float64") = R.multiply(lv1_adjoint, lv) lv_1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv1_adjoint, lv) lv_adjoint1: R.Tensor((2, 2), dtype="float64") = R.add(lv_adjoint, lv_1) x1_adjoint: R.Tensor((2, 2), dtype="float64") = lv_adjoint1 y_adjoint: R.Tensor((2, 2), dtype="float64") = x1_adjoint y_adjoint_out: R.Tensor((2, 2), dtype="float64") = y_adjoint R.output(gv, y_adjoint_out) return (gv, (y_adjoint_out,)) @R.function def optimizer(params: R.Tuple(R.Tensor((2, 2), dtype="float64")), gradients: R.Tuple(R.Tensor((2, 2), dtype="float64")), optim_states: R.Tuple(R.Tensor((), dtype="int64"))) -> R.Tuple(R.Tuple(R.Tensor((2, 2), dtype="float64")), R.Tuple(R.Tensor((), dtype="int64"))): with R.dataflow(): num_steps: R.Tensor((), dtype="int64") = optim_states[0] num_steps_new: R.Tensor((), dtype="int64") = R.add(num_steps, R.const(1, "int64")) y: R.Tensor((2, 2), dtype="float64") = params[0] y_grad: R.Tensor((2, 2), dtype="float64") = gradients[0] lv: R.Tensor((2, 2), dtype="float64") = R.multiply(R.const(0.10000000000000001, "float64"), y_grad) y_new: R.Tensor((2, 2), dtype="float64") = R.subtract(y, lv) params_new: R.Tuple(R.Tensor((2, 2), dtype="float64")) = (y_new,) optim_states_new: R.Tuple(R.Tensor((), dtype="int64")) = (num_steps_new,) R.output(params_new, optim_states_new) return (params_new, optim_states_new) # fmt: on ty = relax.TensorType((2, 2), "float64") setup_trainer = SetupTrainer(MSELoss(reduction="sum"), SGD(0.1), [ty, ty], legalize=False) train_mod = setup_trainer(Backbone) assert_structural_equal(train_mod.without_attr("optim_state"), Expected) def test_states(): # fmt: off @I.ir_module class Backbone: I.module_attrs({"param_num": 1, "state_num": 1}) @R.function def backbone(x: R.Tensor((2, 2), "float64"), y: R.Tensor((2, 2), "float64"), z: R.Tensor((2, 2), "float64")): with R.dataflow(): x1 = x + y z1 = z + R.const(1, "float64") R.output(x1, z1) return x1, z1 @I.ir_module class Expected: I.module_attrs({"input_num": 1, "param_num": 1, "state_num": 1}) @R.function def backbone(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64"), z: R.Tensor((2, 2), dtype="float64")) -> R.Tuple(R.Tensor((2, 2), dtype="float64"), R.Tensor((2, 2), dtype="float64")): with R.dataflow(): x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y) z1: R.Tensor((2, 2), dtype="float64") = R.add(z, R.const(1, "float64")) R.output(x1, z1) return (x1, z1) @R.function def backbone_loss(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64"), z: R.Tensor((2, 2), dtype="float64"), targets: R.Tensor((2, 2), dtype="float64")) -> R.Tuple(R.Tensor((), dtype="float64"), R.Tensor((2, 2), dtype="float64")): with R.dataflow(): x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y) z1: R.Tensor((2, 2), dtype="float64") = R.add(z, R.const(1, "float64")) lv: R.Tensor((2, 2), dtype="float64") = R.subtract(x1, targets) lv1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv, lv) gv: R.Tensor((), dtype="float64") = R.sum(lv1, axis=None, keepdims=False) R.output(z1, gv) return (gv, z1) @R.function def backbone_loss_adjoint(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64"), z: R.Tensor((2, 2), dtype="float64"), targets: R.Tensor((2, 2), dtype="float64")) -> R.Tuple(R.Tuple(R.Tensor((), dtype="float64"), R.Tensor((2, 2), dtype="float64")), R.Tuple(R.Tensor((2, 2), dtype="float64"))): with R.dataflow(): x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y) z1: R.Tensor((2, 2), dtype="float64") = R.add(z, R.const(1, "float64")) lv: R.Tensor((2, 2), dtype="float64") = R.subtract(x1, targets) lv1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv, lv) gv: R.Tensor((), dtype="float64") = R.sum(lv1, axis=None, keepdims=False) gv_adjoint: R.Tensor((), dtype="float64") = R.ones(R.shape([]), dtype="float64") lv1_adjoint: R.Tensor((2, 2), dtype="float64") = R.broadcast_to(gv_adjoint, R.shape([2, 2])) lv_adjoint: R.Tensor((2, 2), dtype="float64") = R.multiply(lv1_adjoint, lv) lv_1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv1_adjoint, lv) lv_adjoint1: R.Tensor((2, 2), dtype="float64") = R.add(lv_adjoint, lv_1) x1_adjoint: R.Tensor((2, 2), dtype="float64") = lv_adjoint1 y_adjoint: R.Tensor((2, 2), dtype="float64") = x1_adjoint y_adjoint_out: R.Tensor((2, 2), dtype="float64") = y_adjoint R.output(z1, gv, y_adjoint_out) return ((gv, z1), (y_adjoint_out,)) @R.function def optimizer(params: R.Tuple(R.Tensor((2, 2), dtype="float64")), gradients: R.Tuple(R.Tensor((2, 2), dtype="float64")), optim_states: R.Tuple(R.Tensor((), dtype="int64"), R.Tensor((2, 2), dtype="float64"))) -> R.Tuple(R.Tuple(R.Tensor((2, 2), dtype="float64")), R.Tuple(R.Tensor((), dtype="int64"), R.Tensor((2, 2), dtype="float64"))): with R.dataflow(): num_steps: R.Tensor((), dtype="int64") = optim_states[0] num_steps_new: R.Tensor((), dtype="int64") = R.add(num_steps, R.const(1, "int64")) y: R.Tensor((2, 2), dtype="float64") = params[0] y_grad: R.Tensor((2, 2), dtype="float64") = gradients[0] y_v: R.Tensor((2, 2), dtype="float64") = optim_states[1] lv: R.Tensor((2, 2), dtype="float64") = R.multiply(R.const(0.10000000000000001, "float64"), y_v) y_v_new: R.Tensor((2, 2), dtype="float64") = R.add(lv, y_grad) lv1: R.Tensor((2, 2), dtype="float64") = R.multiply(R.const(0.10000000000000001, "float64"), y_v_new) y_new: R.Tensor((2, 2), dtype="float64") = R.subtract(y, lv1) params_new: R.Tuple(R.Tensor((2, 2), dtype="float64")) = (y_new,) optim_states_new: R.Tuple(R.Tensor((), dtype="int64"), R.Tensor((2, 2), dtype="float64")) = num_steps_new, y_v_new R.output(params_new, optim_states_new) return (params_new, optim_states_new) # fmt: on ty = relax.TensorType((2, 2), "float64") setup_trainer = SetupTrainer( MSELoss(reduction="sum"), MomentumSGD(0.1, 0.1), [ty, ty], legalize=False ) train_mod = setup_trainer(Backbone) assert_structural_equal(train_mod.without_attr("optim_state"), Expected) def test_invalid_mod(): @I.ir_module class NoAttr: @R.function def backbone( w0: R.Tensor((10, 5), "float32"), b0: R.Tensor((5,), "float32"), x: R.Tensor((1, 10), "float32"), ): with R.dataflow(): lv0 = R.matmul(x, w0) gv = R.add(lv0, b0) out = R.nn.relu(gv) R.output(gv, out) return gv, out pred_ty = relax.TensorType((1, 5), "float32") setup_trainer = SetupTrainer( MSELoss(reduction="sum"), SGD(0.001), [pred_ty, pred_ty], ) with pytest.raises((RuntimeError, ValueError)): SetupTrainer( MSELoss(reduction="sum"), SGD(0.001), [pred_ty, pred_ty], )(NoAttr) @I.ir_module class WrongFuncName: @R.function def main( w0: R.Tensor((10, 5), "float32"), b0: R.Tensor((5,), "float32"), x: R.Tensor((1, 10), "float32"), ): with R.dataflow(): lv0 = R.matmul(x, w0) lv1 = R.add(lv0, b0) out = R.nn.relu(lv1) R.output(out) return out with pytest.raises(ValueError): setup_trainer(WrongFuncName) if __name__ == "__main__": tvm.testing.main()