wehub-resource-sync
74 lines
3.2 KiB
ReStructuredText
74 lines
3.2 KiB
ReStructuredText
.. 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.
|
|
|
|
.. _relax-abstraction:
|
|
|
|
Graph Abstraction for ML Models
|
|
-------------------------------
|
|
Graph abstraction is a key technique used in machine learning (ML) compilers
|
|
to represent and reason about the structure and data flow of ML models. By
|
|
abstracting the model into a graph representation, the compiler can perform
|
|
various optimizations to improve performance and efficiency. This tutorial will
|
|
cover the basics of graph abstraction, its key elements of Relax IR, and how it enables optimization in ML compilers.
|
|
|
|
What is Graph Abstraction?
|
|
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
Graph abstraction is the process of representing an ML model as a directed graph,
|
|
where the nodes represent computational operations (e.g., matrix multiplication,
|
|
convolution) and the edges represent the flow of data between these operations.
|
|
This abstraction allows the compiler to analyze the dependencies and
|
|
relationships between different parts of the model.
|
|
|
|
.. code:: python
|
|
|
|
from tvm.script import relax as R
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 784), dtype="float32"),
|
|
weight: R.Tensor((784, 256), dtype="float32"),
|
|
bias: R.Tensor((256,), dtype="float32"),
|
|
) -> R.Tensor((1, 256), dtype="float32"):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, weight)
|
|
lv1 = R.add(lv0, bias)
|
|
gv = R.nn.relu(lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
Key Features of Relax
|
|
~~~~~~~~~~~~~~~~~~~~~
|
|
Relax, the graph representation utilized in Apache TVM,
|
|
facilitates end-to-end optimization of ML models through several crucial
|
|
features:
|
|
|
|
- **First-class symbolic shape**: Relax employs symbolic shapes to represent
|
|
tensor dimensions, enabling global tracking of dynamic shape relationships
|
|
across tensor operators and function calls.
|
|
|
|
- **Multi-level abstractions**: Relax supports cross-level abstractions, from
|
|
high-level neural network layers to low-level tensor operations, enabling
|
|
optimizations that span different hierarchies within the model.
|
|
|
|
- **Composable transformations**: Relax offers a framework for composable
|
|
transformations that can be selectively applied to different model components.
|
|
This includes capabilities such as partial lowering and partial specialization,
|
|
providing flexible customization and optimization options.
|
|
|
|
These features collectively empower Relax to offer a powerful and adaptable approach
|
|
to ML model optimization within the Apache TVM ecosystem.
|