Ubospica opened a new pull request, #14658: URL: https://github.com/apache/tvm/pull/14658
This PR introduces several utils to support training models in Relax IR, including: - Loss functions. A library that generates relax functions representing common loss functions, such as MSELoss and CrossEntropyLoss. - AppendLoss pass. You can specify one function (the backbone function) in the IRModule, and one loss function to AppendLoss. It will attach the loss function to the backbone function, and add the result function to the IRModule. - Optimizers. A library that generates relax functions representing common optimizers, such as SGD and Adam. - SetupTrainer. A pass that transform the backbone module (a IRModule containing only one function named "backbone") into a module that contains all functions used in the training process. - Trainer. A runtime utility that stores parameters, states, etc, and helps to run the training IRModule. Co-authored-by: Chaofan Lin <[email protected]> -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
