wehub-resource-sync
355 lines
12 KiB
Python
355 lines
12 KiB
Python
# 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.
|
|
# pylint: disable=redefined-builtin, invalid-name
|
|
# ruff: noqa: RUF005
|
|
"""PyTorch-like nn.Module API for constructing workloads."""
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import numpy as np # type: ignore
|
|
|
|
import tvm
|
|
from tvm import relax, tirx, topi
|
|
from tvm.relax.op.grad.grad import end_checkpoint, start_checkpoint
|
|
|
|
|
|
def emit(expr: relax.Expr, name_hint: str = "") -> relax.Var:
|
|
return relax.BlockBuilder.current().emit(expr, name_hint=name_hint)
|
|
|
|
|
|
def emit_te(func: Callable, *args: Any, **kwargs: Any) -> relax.Var:
|
|
return relax.BlockBuilder.current().emit_te(func, *args, **kwargs)
|
|
|
|
|
|
def checkpoint(
|
|
func: Callable, *args: Any, **kwargs: Any
|
|
) -> relax.Var | list[relax.Var] | list[Any]:
|
|
"""Mark function(*args, **kwargs) should be computed in a checkpointed manner during backward.
|
|
|
|
To be specific, args and kwargs will be checkpointed, and func(*args, **kwargs) will be
|
|
recomputed in the backward stage.
|
|
"""
|
|
args = [start_checkpoint(v) if isinstance(v, relax.Expr) else v for v in args]
|
|
kwargs = {k: start_checkpoint(v) if isinstance(v, relax.Expr) else v for k, v in kwargs.items()}
|
|
result = func(*args, **kwargs)
|
|
if isinstance(result, list | tuple):
|
|
result = [end_checkpoint(v) if isinstance(v, relax.Expr) else v for v in result]
|
|
else:
|
|
assert isinstance(result, relax.Expr)
|
|
result = end_checkpoint(result)
|
|
return result
|
|
|
|
|
|
def emit_checkpoint(
|
|
func: Callable, *args: Any, **kwargs: Any
|
|
) -> relax.Var | list[relax.Var] | list[Any]:
|
|
"""Mark function(*args, **kwargs) should be computed in a checkpointed manner during backward.
|
|
|
|
To be specific, args and kwargs will be checkpointed, and func(*args, **kwargs) will be
|
|
recomputed in the backward stage.
|
|
|
|
This interface will additionally emit the exprs marked with start_checkpoint() and
|
|
end_checkpoint() with suffix "_scp" and "_ecp" respectively, for easily understanding the
|
|
result tvmscript.
|
|
"""
|
|
bb = relax.BlockBuilder.current()
|
|
args = [
|
|
bb.emit(start_checkpoint(v), v.name_hint + "_scp") if isinstance(v, relax.Var) else v
|
|
for v in args
|
|
]
|
|
kwargs = {
|
|
k: bb.emit(start_checkpoint(v), v.name_hint + "_scp") if isinstance(v, relax.Var) else v
|
|
for k, v in kwargs.items()
|
|
}
|
|
result = func(*args, **kwargs)
|
|
if isinstance(result, list | tuple):
|
|
result = list(result)
|
|
for i, v in enumerate(result):
|
|
if isinstance(v, relax.Expr):
|
|
if not isinstance(v, relax.Var):
|
|
v = bb.emit(v)
|
|
result[i] = bb.emit(end_checkpoint(v), v.name_hint + "_ecp")
|
|
else:
|
|
assert isinstance(result, relax.Expr)
|
|
result_emit = bb.emit(result)
|
|
result = bb.emit(end_checkpoint(result_emit), result_emit.name_hint + "_ecp")
|
|
|
|
return result
|
|
|
|
|
|
def emit_checkpoint_sequential(
|
|
functions: list[Callable],
|
|
segments: int | list[int],
|
|
input: relax.Var,
|
|
checkpoint_last: bool = False,
|
|
) -> relax.Var | list[relax.Var] | list[Any]:
|
|
"""A helper function for checkpointing sequential models. This interface has similar purpose
|
|
as torch.utils.checkpoint.checkpoint_sequential.
|
|
|
|
Sequential models execute a list of modules/functions in order (sequentially). Therefore, we
|
|
can divide such a model in various segments and checkpoint each segment. By default, we will
|
|
checkpoint all segments except the last, meaning their inputs will be saved from the forward
|
|
stage and they will be recomputed in the backward stage.
|
|
|
|
Parameters
|
|
----------
|
|
functions : List[Callable]
|
|
The list of functions to be executed sequentially.
|
|
|
|
segments : Union[int, List[int]]
|
|
The segments. If segments is int `n`, functions will be evenly divided into `n` segments;
|
|
if segments is a list of ints, it marks the start of every segment.
|
|
|
|
input : relax.Var
|
|
The input of the first function.
|
|
|
|
checkpoint_last : bool
|
|
Whether the last segment will be checkpointed. Default: False
|
|
|
|
Returns
|
|
-------
|
|
output : Union[relax.Var, List[relax.Var], List[Any]]
|
|
The emited output of the last function.
|
|
"""
|
|
bb = relax.BlockBuilder.current()
|
|
|
|
def run_function(start, end, functions):
|
|
def forward(input):
|
|
for j in range(start, end):
|
|
input = functions[j](input)
|
|
return input
|
|
|
|
return forward
|
|
|
|
n = len(functions)
|
|
if not isinstance(segments, list):
|
|
segments = list(range(0, n, n // segments)) + [n]
|
|
if segments[-1] != n:
|
|
segments = segments + [n]
|
|
|
|
assert len(segments) >= 2
|
|
|
|
for i in range(len(segments) - 1):
|
|
if i == len(segments) - 2 and not checkpoint_last:
|
|
input = run_function(segments[i], segments[i + 1], functions)(input)
|
|
else:
|
|
input = emit_checkpoint(run_function(segments[i], segments[i + 1], functions), input)
|
|
|
|
assert isinstance(input, relax.Expr)
|
|
if not isinstance(input, relax.Var):
|
|
input = bb.emit(input)
|
|
return input
|
|
|
|
|
|
def _try_unique_name(name: str):
|
|
"""Attempt to uniquify the name
|
|
|
|
If a `relax.BlockBuilder` is active, use it to return a unique
|
|
name. Otherwise, return the name itself.
|
|
|
|
Two distinct variables in Relax may have identical names.
|
|
However, for user readability, it is convenient to have all names
|
|
be unique within a Relax function. If a Placeholder or Parameter
|
|
is defined within an active `relax.BlockBuilder`, that context may
|
|
be used to provide a unique name. Otherwise, allow the duplicate
|
|
names.
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
|
|
The variable name
|
|
|
|
Returns
|
|
-------
|
|
updated_name: str
|
|
|
|
The updated variable name
|
|
|
|
|
|
"""
|
|
block_builder = relax.BlockBuilder.current()
|
|
if block_builder is None:
|
|
return name
|
|
else:
|
|
return block_builder.get_unique_name(name)
|
|
|
|
|
|
class Placeholder(relax.Var):
|
|
"""A placeholder variable that can represent model input."""
|
|
|
|
def __init__(self, shape: list[Any] | tuple[Any, ...], dtype="float32", name="data"):
|
|
if not isinstance(shape, list | tuple):
|
|
raise TypeError("the shape of Placeholder is expected to be a list or a tuple")
|
|
super().__init__(_try_unique_name(name), relax.TensorType(shape, dtype))
|
|
|
|
|
|
class Parameter(relax.Var):
|
|
"""A special kind of relax Var that represents model parameter(weight)."""
|
|
|
|
def __init__(self, shape: list[Any] | tuple[Any, ...], dtype="float32", name="param"):
|
|
if not isinstance(shape, list | tuple):
|
|
raise TypeError("the shape of Parameter is expected to be a list or a tuple")
|
|
|
|
super().__init__(_try_unique_name(name), relax.TensorType(shape, dtype))
|
|
|
|
|
|
class Module(tvm.relax.frontend.nn.SubroutineMixin):
|
|
"""Base class for all model modules.
|
|
|
|
A neural network or a layer can subclass this class.
|
|
|
|
By default, calls into this module will generate the `relax.Expr`
|
|
representing the output within the current function body. Setting
|
|
the variable "define_subrouine" to True; either at the
|
|
`nn.Module`, subclass, or instance level; will instead produce a
|
|
subroutine within the same module, which is then called within the
|
|
current function body.
|
|
|
|
Example
|
|
-------
|
|
.. code-block:: python
|
|
|
|
# Define a linear layer
|
|
class Linear(Module)
|
|
def __init__(self, in_features, out_features, bias=True):
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.weight = Parameter((in_features, out_features), name="linear_weight")
|
|
if bias:
|
|
self.bias = Parameter((out_features,), name="linear_bias")
|
|
else:
|
|
self.bias = None
|
|
|
|
# All submodules should implement forward.
|
|
# Defines the forward computation performed at every call.
|
|
def forward(self, input: relax.Expr) -> relax.Var:
|
|
y = emit_te(topi.matmul, input, self.weight)
|
|
if self.bias is not None:
|
|
y = emit_te(topi.add, y, self.bias)
|
|
return y
|
|
"""
|
|
|
|
define_subroutine: bool = False
|
|
|
|
def parameters(self) -> list[Parameter]:
|
|
"""Return the list of parameters in the module."""
|
|
return _unpack_params(self.__dict__)
|
|
|
|
def forward(self, input: relax.Expr):
|
|
"""Define the computation performed at every call."""
|
|
raise NotImplementedError()
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return self.forward(*args, **kwargs)
|
|
|
|
|
|
def _unpack_params(value: object) -> list[relax.Var]:
|
|
if isinstance(value, Parameter):
|
|
return [value]
|
|
if isinstance(value, Module):
|
|
return value.parameters()
|
|
if isinstance(value, dict):
|
|
params = []
|
|
for v in value.values():
|
|
params += _unpack_params(v)
|
|
return params
|
|
if isinstance(value, list | tuple):
|
|
params = []
|
|
for v in value:
|
|
params += _unpack_params(v)
|
|
return params
|
|
return []
|
|
|
|
|
|
def init_params(mod: tvm.IRModule) -> list[tvm.runtime.Tensor]:
|
|
"""Utility function to initialize model's parameters."""
|
|
shape_dict = {v.name_hint: v.ty.shape for v in mod["main"].params}
|
|
params = []
|
|
for k, v in shape_dict.items():
|
|
if k.startswith("data"):
|
|
continue
|
|
if isinstance(v, relax.ShapeExpr):
|
|
shape = []
|
|
for i in v:
|
|
if isinstance(i, tirx.IntImm):
|
|
shape.append(int(i))
|
|
else:
|
|
raise TypeError("cannot initialize for unknown-shape parameters.")
|
|
params.append(tvm.runtime.tensor(np.zeros(shape).astype(np.float32)))
|
|
else:
|
|
raise TypeError("cannot initialize for unknown-shape parameters.")
|
|
return params
|
|
|
|
|
|
class Sequential(Module):
|
|
"""A sequential container that concatenates modules in it.
|
|
|
|
Example
|
|
-------
|
|
.. code-block:: python
|
|
|
|
model = nn.Sequential(
|
|
nn.Conv2d(1, 20, 5),
|
|
nn.ReLU(),
|
|
nn.Conv2d(20, 64, 5),
|
|
nn.ReLU()
|
|
)
|
|
"""
|
|
|
|
def __init__(self, *modules: Module):
|
|
self.modules = modules
|
|
|
|
def forward(self, input: relax.Expr) -> relax.Var:
|
|
for module in self.modules:
|
|
input = module(input)
|
|
return input
|
|
|
|
|
|
class ReLU(Module):
|
|
"""Applies the rectified linear unit activation function on the input."""
|
|
|
|
def forward(self, input: relax.Expr) -> relax.Var:
|
|
return emit_te(topi.nn.relu, input)
|
|
|
|
|
|
class LogSoftmax(Module):
|
|
"""Applies log softmax activation function on the input."""
|
|
|
|
def forward(self, input: relax.Expr) -> relax.Var:
|
|
return emit_te(topi.nn.log_softmax, input)
|
|
|
|
|
|
class Linear(Module):
|
|
"""Applies a linear transformation to the input data: :math:`y = xA + b`."""
|
|
|
|
def __init__(self, in_features, out_features, bias=True):
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.weight = Parameter((in_features, out_features), name="linear_weight")
|
|
if bias:
|
|
self.bias = Parameter((out_features,), name="linear_bias")
|
|
else:
|
|
self.bias = None
|
|
|
|
def forward(self, input: relax.Expr) -> relax.Var:
|
|
y = emit_te(topi.matmul, input, self.weight)
|
|
if self.bias is not None:
|
|
y = emit_te(topi.add, y, self.bias)
|
|
return y
|