chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,354 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user