wehub-resource-sync
217 lines
6.1 KiB
Python
217 lines
6.1 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
|
|
"""Operators to implement operaor gradients. Used in `_op_gradient.py`.
|
|
|
|
We are trying to keep grad operators as simple as possible, and hope they are only used for finding
|
|
gradients for forward operators. The ty inference for grad operators just returns the
|
|
ty of the input.
|
|
"""
|
|
|
|
from ...expr import Expr
|
|
from . import _ffi_api
|
|
|
|
|
|
def no_grad(input: Expr) -> Expr:
|
|
"""No gradient dummy operator w.r.t. the input.
|
|
|
|
Parameters
|
|
----------
|
|
input : relax.Expr
|
|
The corresponding input tensor.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The no-gradient representation w.r.t. input.
|
|
"""
|
|
return _ffi_api.no_grad(input) # type: ignore
|
|
|
|
|
|
def start_checkpoint(input: Expr) -> Expr:
|
|
"""Mark the start of the checkpoint stage. The computation between start_checkpoint and
|
|
end_checkpoint will be marked as the checkpoint stage.
|
|
|
|
Rather than storing all intermediate activations of the entire computation graph for
|
|
computing backward, the checkpointed stage does not save intermediate activations, and instead
|
|
recomputes them in backward process.
|
|
|
|
For instance,
|
|
```
|
|
a = relax.Var("a", relax.TensorType((2, 2), "float32"))
|
|
b = relax.Var("b", relax.TensorType((2, 2), "float32"))
|
|
c = a * 2
|
|
d = b * 2
|
|
c_cp = start_checkpoint(c)
|
|
d_cp = start_checkpoint(d)
|
|
e = c_cp + d_cp
|
|
e_out = end_checkpoint(e)
|
|
```
|
|
Then `e` will be recomputed in the backward stage.
|
|
|
|
See tvm.relax.transform.Gradient, tvm.relax.testing.nn.checkpoint,
|
|
tvm.relax.op.grad.end_checkpoint for more information.
|
|
|
|
Parameters
|
|
----------
|
|
input : relax.Expr
|
|
The tensor marking the input of the checkpoint stage.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The same tensor as the input.
|
|
"""
|
|
return _ffi_api.start_checkpoint(input) # type: ignore
|
|
|
|
|
|
def end_checkpoint(input: Expr) -> Expr:
|
|
"""Mark the end of checkpoint stage. See tvm.relax.op.grad.start_checkpoint.
|
|
|
|
Parameters
|
|
----------
|
|
input : relax.Expr
|
|
The output of the checkpoint stage.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The same tensor as the input.
|
|
"""
|
|
return _ffi_api.end_checkpoint(input) # type: ignore
|
|
|
|
|
|
def nll_loss_backward(
|
|
output_grad: Expr,
|
|
predictions: Expr,
|
|
targets: Expr,
|
|
weights: Expr | None = None,
|
|
reduction: str = "mean",
|
|
ignore_index: int = -100,
|
|
) -> Expr:
|
|
"""Backward operator of relax.nn.nll_loss. All parameters except output_grad is the same as
|
|
relax.nn.nll_loss. Returns the gradient w.r.t. predictions.
|
|
|
|
Parameters
|
|
----------
|
|
output_grad : relax.Expr
|
|
The gradient w.r.t. the result of nll_loss.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The gradient w.r.t. predictions.
|
|
"""
|
|
return _ffi_api.nll_loss_backward( # type: ignore
|
|
output_grad, predictions, targets, weights, reduction, ignore_index
|
|
)
|
|
|
|
|
|
def max_pool2d_backward(
|
|
output_grad: Expr,
|
|
data: Expr,
|
|
pool_size: tuple[int, int] = (1, 1),
|
|
strides: tuple[int, int] = (1, 1),
|
|
padding: tuple[int, int, int, int] = (0, 0, 0, 0),
|
|
dilation: tuple[int, int] = (1, 1),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = False,
|
|
layout: str = "NCHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
"""Backward operator of relax.nn.max_pool2d. All parameters except output_grad is the same as
|
|
relax.nn.max_pool2d. Returns the gradient w.r.t. data.
|
|
|
|
Parameters
|
|
----------
|
|
output_grad : relax.Expr
|
|
The gradient w.r.t. the result of max_pool2d.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The gradient w.r.t. data.
|
|
"""
|
|
return _ffi_api.max_pool2d_backward( # type: ignore
|
|
output_grad,
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def avg_pool2d_backward(
|
|
output_grad: Expr,
|
|
data: Expr,
|
|
pool_size: tuple[int, int] = (1, 1),
|
|
strides: tuple[int, int] = (1, 1),
|
|
padding: tuple[int, int, int, int] = (0, 0, 0, 0),
|
|
dilation: tuple[int, int] = (1, 1),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = False,
|
|
layout: str = "NCHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
"""Backward operator of relax.nn.avg_pool2d. All parameters except output_grad is the same as
|
|
relax.nn.avg_pool2d. Returns the gradient w.r.t. data.
|
|
|
|
Parameters
|
|
----------
|
|
output_grad : relax.Expr
|
|
The gradient w.r.t. the result of avg_pool2d.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The gradient w.r.t. data.
|
|
"""
|
|
return _ffi_api.avg_pool2d_backward( # type: ignore
|
|
output_grad,
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def take_backward(output_grad: Expr, x: Expr, indices: Expr, axis: int | None = None) -> Expr:
|
|
"""Backward operator of relax.take. All parameters except output_grad is the same as
|
|
relax.take. Returns the gradient w.r.t. x.
|
|
|
|
Parameters
|
|
----------
|
|
output_grad : relax.Expr
|
|
The gradient w.r.t. the result of take.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The gradient w.r.t. x.
|
|
"""
|
|
return _ffi_api.take_backward(output_grad, x, indices, axis) # type: ignore
|