wehub-resource-sync
1283 lines
50 KiB
Python
1283 lines
50 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=invalid-name, inconsistent-return-statements, unidiomatic-typecheck
|
|
# pylint: disable=import-outside-toplevel
|
|
"""PyTorch FX frontend of Relax."""
|
|
|
|
from collections.abc import Callable
|
|
from functools import partial, reduce
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
|
|
from .base_fx_graph_translator import BaseFXGraphImporter
|
|
|
|
|
|
class TorchFXImporter(BaseFXGraphImporter):
|
|
"""An importer from PyTorch FX to Relax."""
|
|
|
|
import torch # type: ignore
|
|
from torch import fx
|
|
|
|
def __init__(self, default_image_layout: str = "NCHW") -> None:
|
|
import torch # type: ignore
|
|
|
|
super().__init__()
|
|
self.named_modules: dict[str, torch.Module] = None
|
|
self.default_image_layout = default_image_layout
|
|
|
|
########## Utilities ##########
|
|
|
|
def _fetch_attr(self, model, target: str):
|
|
import torch # type: ignore
|
|
|
|
target_atoms = target.split(".")
|
|
attr_itr = model
|
|
for i, atom in enumerate(target_atoms):
|
|
if not hasattr(attr_itr, atom):
|
|
raise RuntimeError(
|
|
f"Node referenced non existing target {'.'.join(target_atoms[:i])}"
|
|
)
|
|
attr_itr = getattr(attr_itr, atom)
|
|
if isinstance(attr_itr, torch.Tensor):
|
|
# Its possible for the resulting tensor to be a parameter.
|
|
# If so, return the parameter instead.
|
|
if attr_itr in self.params:
|
|
return self.params[attr_itr]
|
|
return self._convert_torch_tensor_to_relax(attr_itr)
|
|
return attr_itr
|
|
|
|
########## Unary Ops ##########
|
|
|
|
def _reciprocal(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.ty.dtype), x))
|
|
|
|
def _leakyrelu_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
alpha = module.negative_slope
|
|
return self.block_builder.emit(relax.op.nn.leakyrelu(x, alpha))
|
|
|
|
def _softplus_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
beta = module.beta
|
|
threshold = module.threshold
|
|
return self.block_builder.emit(relax.op.nn.softplus(x, beta, threshold))
|
|
|
|
def _log2(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
return self.block_builder.emit(
|
|
relax.op.divide(relax.op.log(x), relax.const(0.6931471805599453, x.ty.dtype))
|
|
)
|
|
|
|
def _log10(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
return self.block_builder.emit(
|
|
relax.op.divide(relax.op.log(x), relax.const(2.302585092994046, x.ty.dtype))
|
|
)
|
|
|
|
def _log1p(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
one = relax.const(1, x.ty.dtype)
|
|
return self.block_builder.emit(relax.op.log(relax.op.add(x, one)))
|
|
|
|
def _sqrt(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
dtype = x.ty.dtype
|
|
|
|
# Check if input is integer type and convert to float32 if needed
|
|
if dtype in ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]:
|
|
x = self.block_builder.emit(relax.op.astype(x, "float32"))
|
|
|
|
return self.block_builder.emit(relax.op.sqrt(x))
|
|
|
|
def _rsqrt(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
dtype = x.ty.dtype
|
|
|
|
# Check if input is integer type and convert to float32 if needed
|
|
if dtype in ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]:
|
|
x = self.block_builder.emit(relax.op.astype(x, "float32"))
|
|
|
|
return self.block_builder.emit(relax.op.rsqrt(x))
|
|
|
|
def _log_softmax_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
dim = module.dim
|
|
assert dim is not None
|
|
return self.block_builder.emit(relax.op.nn.log_softmax(x, dim))
|
|
|
|
def _prelu_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
alpha_tensor = module.weight.numpy()
|
|
alpha = relax.const(alpha_tensor, dtype="float32")
|
|
axis = 0 if len(x.ty.shape) == 1 else 1 # Extract Channel size
|
|
return self.block_builder.emit(relax.op.nn.prelu(x, alpha, axis))
|
|
|
|
def _softmax_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
dim = module.dim
|
|
assert dim is not None
|
|
return self.block_builder.emit(relax.op.nn.softmax(x, dim))
|
|
|
|
def _inplace_tril_triu(self, op: Callable) -> Callable:
|
|
from torch import fx
|
|
|
|
def convert(node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
k = node.args[1] if len(node.args) > 1 else 0
|
|
assert isinstance(k, int)
|
|
|
|
mutated = self.block_builder.emit(op(x, k))
|
|
self.env[node.args[0]] = mutated
|
|
return mutated
|
|
|
|
return convert
|
|
|
|
########## Binary Ops ##############
|
|
|
|
def _binary_op_inplace(self, relax_op: Callable, intrinsic_op: Callable) -> Callable:
|
|
from torch import fx
|
|
|
|
def convert(node: fx.Node) -> relax.Var:
|
|
def promote_binary_op_args(lhs, rhs):
|
|
if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
|
|
return lhs, rhs
|
|
elif isinstance(lhs, relax.Expr):
|
|
assert isinstance(lhs.ty, relax.TensorType)
|
|
return lhs, relax.const(rhs, lhs.ty.dtype)
|
|
elif isinstance(rhs, relax.Expr):
|
|
assert isinstance(rhs.ty, relax.TensorType)
|
|
return relax.const(lhs, rhs.ty.dtype), rhs
|
|
else:
|
|
assert False
|
|
|
|
def call_binary_op(op, lhs, rhs):
|
|
lhs, rhs = promote_binary_op_args(lhs, rhs)
|
|
return self.block_builder.emit(op(lhs, rhs))
|
|
|
|
lhs, rhs = self.retrieve_args(node)
|
|
if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
|
|
output = call_binary_op(relax_op, lhs, rhs)
|
|
self.env[node.args[0]] = output
|
|
return output
|
|
|
|
elif isinstance(lhs, relax.expr.Constant):
|
|
output = call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype))
|
|
self.env[node.args[0]] = output
|
|
return output
|
|
|
|
elif isinstance(rhs, relax.expr.Constant):
|
|
output = call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs)
|
|
self.env[node.args[0]] = output
|
|
return output
|
|
|
|
output = intrinsic_op(lhs, rhs)
|
|
self.env[node.args[0]] = output
|
|
return output
|
|
|
|
return convert
|
|
|
|
########## Neural Network ##########
|
|
|
|
def _adaptive_avg_pool1d_module(self, node: fx.Node) -> relax.Var:
|
|
module = self.named_modules[node.target]
|
|
x = self.env[node.args[0]]
|
|
output_size = module.output_size
|
|
# Expand to 3D by adding batch dim if input is 2D
|
|
x_ndim = x.ty.ndim
|
|
if x_ndim == 2:
|
|
x = relax.op.expand_dims(x, axis=0)
|
|
result = self.block_builder.emit(
|
|
relax.op.nn.adaptive_avg_pool1d(x, output_size, layout="NCW") # (N, C, L)
|
|
)
|
|
# Remove added batch dim from result
|
|
if x_ndim == 2:
|
|
result = relax.op.squeeze(result, axis=[0])
|
|
return result
|
|
|
|
def _adaptive_avg_pool2d_module(self, node: fx.Node) -> relax.Var:
|
|
module = self.named_modules[node.target]
|
|
x = self.env[node.args[0]]
|
|
output_size = module.output_size
|
|
# Expand to 4D by adding batch dim if input is 3D
|
|
x_ndim = x.ty.ndim
|
|
if x_ndim == 3:
|
|
x = relax.op.expand_dims(x, axis=0)
|
|
result = self.block_builder.emit(
|
|
relax.op.nn.adaptive_avg_pool2d(x, output_size, layout="NCHW")
|
|
)
|
|
# Remove added batch dim from result
|
|
if x_ndim == 3:
|
|
result = relax.op.squeeze(result, axis=[0])
|
|
return result
|
|
|
|
def _adaptive_avg_pool3d_module(self, node: fx.Node) -> relax.Var:
|
|
module = self.named_modules[node.target]
|
|
x = self.env[node.args[0]]
|
|
output_size = module.output_size
|
|
# Expand to 5D by adding batch dim if input is 4D
|
|
x_ndim = x.ty.ndim
|
|
if x_ndim == 4:
|
|
x = relax.op.expand_dims(x, axis=0)
|
|
result = self.block_builder.emit(
|
|
relax.op.nn.adaptive_avg_pool3d(x, output_size, layout="NCDHW") # (N, C, D, H, W)
|
|
)
|
|
# Remove added batch dim from result
|
|
if x_ndim == 4:
|
|
result = relax.op.squeeze(result, axis=[0])
|
|
return result
|
|
|
|
def _avg_pool1d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
kernel_size = module.kernel_size
|
|
stride = module.stride
|
|
padding = module.padding
|
|
ceil_mode = module.ceil_mode
|
|
return self._avg_pool1d_impl(x, kernel_size, stride, padding, ceil_mode)
|
|
|
|
def _avg_pool2d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
kernel_size = module.kernel_size
|
|
stride = module.stride
|
|
padding = module.padding
|
|
ceil_mode = module.ceil_mode
|
|
return self._avg_pool2d_impl(x, kernel_size, stride, padding, ceil_mode)
|
|
|
|
def _avg_pool3d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
kernel_size = module.kernel_size
|
|
stride = module.stride
|
|
padding = module.padding
|
|
ceil_mode = module.ceil_mode
|
|
return self._avg_pool3d_impl(x, kernel_size, stride, padding, ceil_mode)
|
|
|
|
def _batch_norm_2d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
bias = self.params[module.bias]
|
|
running_mean = self._convert_torch_tensor_to_relax(module.running_mean)
|
|
running_var = self._convert_torch_tensor_to_relax(module.running_var)
|
|
eps = module.eps
|
|
|
|
res_tuple = self.block_builder.emit(
|
|
relax.op.nn.batch_norm(
|
|
x,
|
|
weight,
|
|
bias,
|
|
running_mean,
|
|
running_var,
|
|
axis=1,
|
|
epsilon=eps,
|
|
)
|
|
)
|
|
|
|
return self.block_builder.emit(relax.TupleGetItem(res_tuple, 0))
|
|
|
|
def _instance_norm(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
|
|
if module.affine:
|
|
weight = self.params[module.weight]
|
|
bias = self.params[module.bias]
|
|
else:
|
|
import numpy as np
|
|
|
|
dtype = x.ty.dtype
|
|
channel = int(self.shape_of(x)[1])
|
|
weight = relax.const(np.ones(channel), dtype=dtype)
|
|
bias = relax.const(np.zeros(channel), dtype=dtype)
|
|
|
|
eps = module.eps
|
|
channel_axis = 1
|
|
dim = len(self.shape_of(x))
|
|
|
|
return self.block_builder.emit(
|
|
relax.op.nn.instance_norm(
|
|
x,
|
|
weight,
|
|
bias,
|
|
channel_axis=channel_axis,
|
|
axes=list(range(2, dim)),
|
|
epsilon=eps,
|
|
)
|
|
)
|
|
|
|
def _conv_transpose1d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
bias = self.params.get(module.bias, None)
|
|
|
|
return self._conv_transpose1d_impl(
|
|
x,
|
|
weight,
|
|
bias=bias,
|
|
strides=module.stride,
|
|
padding=module.padding,
|
|
dilation=module.dilation,
|
|
groups=module.groups,
|
|
output_padding=module.output_padding,
|
|
)
|
|
|
|
def _conv_transpose2d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
bias = self.params.get(module.bias, None)
|
|
|
|
return self._conv_transpose2d_impl(
|
|
x,
|
|
weight,
|
|
bias=bias,
|
|
strides=module.stride,
|
|
padding=module.padding,
|
|
dilation=module.dilation,
|
|
groups=module.groups,
|
|
output_padding=module.output_padding,
|
|
)
|
|
|
|
def _conv1d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
bias = self.params.get(module.bias, None)
|
|
|
|
return self._conv1d_impl(
|
|
x,
|
|
weight,
|
|
bias=bias,
|
|
strides=module.stride,
|
|
padding=module.padding,
|
|
dilation=module.dilation,
|
|
groups=module.groups,
|
|
)
|
|
|
|
def _conv2d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
bias = self.params.get(module.bias, None)
|
|
|
|
return self._conv2d_impl(
|
|
x,
|
|
weight,
|
|
bias=bias,
|
|
strides=module.stride,
|
|
padding=module.padding,
|
|
dilation=module.dilation,
|
|
groups=module.groups,
|
|
)
|
|
|
|
def _conv3d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
bias = self.params.get(module.bias, None)
|
|
|
|
return self._conv3d_impl(
|
|
x,
|
|
weight,
|
|
bias=bias,
|
|
strides=module.stride,
|
|
padding=module.padding,
|
|
dilation=module.dilation,
|
|
groups=module.groups,
|
|
)
|
|
|
|
def _cross_entropy(self, node: fx.Node) -> relax.Expr:
|
|
preds = self.env[node.args[0]]
|
|
targets = self.env[node.args[1]]
|
|
weights = self.env.get(node.kwargs["weight"], None)
|
|
reduction = node.kwargs["reduction"]
|
|
ignore_index = node.kwargs["ignore_index"]
|
|
return self._cross_entropy_loss(preds, targets, weights, reduction, ignore_index)
|
|
|
|
def _cross_entropy_module(self, node: fx.Node) -> relax.Expr:
|
|
preds = self.env[node.args[0]]
|
|
targets = self.env[node.args[1]]
|
|
module = self.named_modules[node.target]
|
|
|
|
weights = module.weight
|
|
if weights is not None:
|
|
if weights in self.params:
|
|
weights = self.params[weights]
|
|
else:
|
|
weights = relax.const(weights.numpy(), preds.ty.dtype)
|
|
|
|
reduction = module.reduction
|
|
ignore_index = module.ignore_index
|
|
|
|
return self._cross_entropy_loss(
|
|
preds,
|
|
targets,
|
|
weights,
|
|
reduction,
|
|
ignore_index,
|
|
)
|
|
|
|
def _embedding_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
return self._embedding_impl(x, weight)
|
|
|
|
def _group_norm_module(self, node: fx.Node) -> relax.Var:
|
|
import torch # type: ignore
|
|
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
num_groups = module.num_groups
|
|
if module.affine:
|
|
gamma = self.params[module.weight]
|
|
beta = self.params[module.bias]
|
|
else:
|
|
gamma = relax.const(torch.ones_like(module.num_channels), x.ty.dtype)
|
|
beta = relax.const(torch.zeros_like(module.num_channels), x.ty.dtype)
|
|
eps = module.eps
|
|
|
|
dim = len(self.shape_of(x))
|
|
return self.block_builder.emit(
|
|
relax.op.nn.group_norm(
|
|
x,
|
|
gamma,
|
|
beta,
|
|
num_groups=num_groups,
|
|
channel_axis=1,
|
|
axes=list(range(2, dim)),
|
|
epsilon=eps,
|
|
)
|
|
)
|
|
|
|
def _interpolate(self, node: fx.Node) -> relax.Var:
|
|
# torch.nn.functional.interpolate(
|
|
# input, size=None, scale_factor=None, mode='nearest', align_corners=None,
|
|
# recompute_scale_factor=None, antialias=False)
|
|
data = self.env[node.args[0]]
|
|
size = (
|
|
node.args[1]
|
|
if len(node.args) > 1
|
|
else (node.kwargs["size"] if "size" in node.kwargs else None)
|
|
)
|
|
scale_factor = (
|
|
node.args[2]
|
|
if len(node.args) > 2
|
|
else (node.kwargs["scale_factor"] if "scale_factor" in node.kwargs else None)
|
|
)
|
|
method = (
|
|
node.args[3]
|
|
if len(node.args) > 3
|
|
else (node.kwargs["mode"] if "mode" in node.kwargs else "nearest")
|
|
)
|
|
align_corners = (
|
|
node.args[4]
|
|
if len(node.args) > 4
|
|
else (node.kwargs["align_corners"] if "align_corners" in node.kwargs else None)
|
|
)
|
|
recompute_scale_factor = (
|
|
node.args[5]
|
|
if len(node.args) > 5
|
|
else (
|
|
node.kwargs["recompute_scale_factor"]
|
|
if "recompute_scale_factor" in node.kwargs
|
|
else None
|
|
)
|
|
)
|
|
antialias = (
|
|
node.args[6]
|
|
if len(node.args) > 6
|
|
else (node.kwargs["antialias"] if "antialias" in node.kwargs else False)
|
|
)
|
|
|
|
assert recompute_scale_factor is None
|
|
assert antialias is False
|
|
|
|
if size is None:
|
|
shape = self.shape_of(data)
|
|
assert isinstance(shape, relax.ShapeExpr)
|
|
# Determine spatial dimension indices based on layout
|
|
# NCHW: spatial dims are [2, 3, ...] (skip batch and channel)
|
|
# NHWC: spatial dims are [1, 2, ...] (skip batch, before channel)
|
|
if self.default_image_layout in ("NHWC", "NDHWC"):
|
|
spatial_start = 1
|
|
spatial_end = len(shape) - 1
|
|
else: # NCHW or other layouts
|
|
spatial_start = 2
|
|
spatial_end = len(shape)
|
|
|
|
if isinstance(scale_factor, tuple):
|
|
assert len(scale_factor) == spatial_end - spatial_start
|
|
size = tuple(
|
|
int(shape[i].value * scale_factor[i - spatial_start])
|
|
for i in range(spatial_start, spatial_end)
|
|
)
|
|
else:
|
|
size = tuple(
|
|
int(shape[i].value * scale_factor) for i in range(spatial_start, spatial_end)
|
|
)
|
|
|
|
if method.startswith("nearest"):
|
|
method = "nearest_neighbor"
|
|
elif method.startswith("bi"):
|
|
method = method[2:]
|
|
elif method.startswith("tri"):
|
|
method = method[3:]
|
|
|
|
if method == "nearest_neighbor":
|
|
coord_trans = "asymmetric"
|
|
elif align_corners is True:
|
|
coord_trans = "align_corners"
|
|
else:
|
|
coord_trans = "half_pixel"
|
|
|
|
if data.ty.ndim == 5:
|
|
if self.default_image_layout == "NDHWC":
|
|
layout_3d = "NDHWC"
|
|
else:
|
|
layout_3d = "NCDHW"
|
|
|
|
return self.block_builder.emit(
|
|
relax.op.image.resize3d(
|
|
data,
|
|
size,
|
|
layout=layout_3d,
|
|
method=method,
|
|
coordinate_transformation_mode=coord_trans,
|
|
)
|
|
)
|
|
else:
|
|
return self.block_builder.emit(
|
|
relax.op.image.resize2d(
|
|
data,
|
|
size,
|
|
layout=self.default_image_layout,
|
|
method=method,
|
|
coordinate_transformation_mode=coord_trans,
|
|
)
|
|
)
|
|
|
|
def _linear_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
weight = self.params[module.weight]
|
|
bias = self.params.get(module.bias, None)
|
|
return self.block_builder.emit(relax.op.linear(x, weight, bias, "float32"))
|
|
|
|
def _max_pool1d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
kernel_size = module.kernel_size
|
|
stride = module.stride
|
|
padding = module.padding
|
|
dilation = module.dilation
|
|
ceil_mode = module.ceil_mode
|
|
|
|
return self._max_pool1d_impl(x, kernel_size, stride, padding, dilation, ceil_mode)
|
|
|
|
def _max_pool2d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
kernel_size = module.kernel_size
|
|
stride = module.stride
|
|
padding = module.padding
|
|
dilation = module.dilation
|
|
ceil_mode = module.ceil_mode
|
|
|
|
return self._max_pool2d_impl(x, kernel_size, stride, padding, dilation, ceil_mode)
|
|
|
|
def _max_pool3d_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
kernel_size = module.kernel_size
|
|
stride = module.stride
|
|
padding = module.padding
|
|
dilation = module.dilation
|
|
ceil_mode = module.ceil_mode
|
|
|
|
return self._max_pool3d_impl(x, kernel_size, stride, padding, dilation, ceil_mode)
|
|
|
|
def _pixel_shuffle_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
upscale_factor = module.upscale_factor
|
|
|
|
return self.block_builder.emit(relax.op.nn.pixel_shuffle(x, upscale_factor))
|
|
|
|
########## Linear Interpolation ##########
|
|
|
|
def _lerp(self, node: fx.Node) -> relax.Var:
|
|
start = self.env[node.args[0]]
|
|
end = self.env[node.args[1]]
|
|
weight = self.env[node.args[2]]
|
|
return self.block_builder.emit(
|
|
relax.op.add(start, relax.op.multiply(weight, relax.op.subtract(end, start)))
|
|
)
|
|
|
|
########## Manipulation ##########
|
|
|
|
def _chunk(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
chunks = node.args[1]
|
|
dim = node.args[2] if len(node.args) > 2 else node.kwargs.get("dim", 0)
|
|
return self.block_builder.emit(relax.op.split(x, chunks, dim))
|
|
|
|
def _flatten_module(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
module = self.named_modules[node.target]
|
|
start_dim = module.start_dim
|
|
end_dim = module.end_dim
|
|
return self._flatten_impl(x, start_dim, end_dim)
|
|
|
|
def _narrow(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
dim = node.args[1]
|
|
start = node.args[2]
|
|
length = node.args[3]
|
|
return self.block_builder.emit(relax.op.strided_slice(x, [dim], [start], [length]))
|
|
|
|
def _numel(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
shape = self.shape_of(x)
|
|
return relax.const(reduce(lambda x, y: x * y, [s.value for s in shape]), "int32")
|
|
|
|
def _select(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
dim = node.args[1]
|
|
index = relax.const(node.args[2], "int64")
|
|
return self.block_builder.emit(relax.op.take(x, index, dim))
|
|
|
|
def _size(self, node: fx.Node) -> relax.Expr:
|
|
x = self.env[node.args[0]]
|
|
shape = self.shape_of(x)
|
|
if len(node.args) == 1:
|
|
assert isinstance(shape, relax.ShapeExpr)
|
|
return shape
|
|
assert len(node.args) == 2
|
|
idx = node.args[1]
|
|
return self.shape_of(x)[idx].value
|
|
|
|
########## Creation ##########
|
|
|
|
def _inplace_copy(self, node: fx.Node) -> relax.Var:
|
|
dest = self.env[node.args[0]]
|
|
src = self.env[node.args[1]]
|
|
|
|
if src.ty.dtype != dest.ty.dtype:
|
|
src = self.block_builder.emit(relax.op.astype(src, dest.ty.dtype.dtype))
|
|
|
|
dest_shape = self.shape_of(dest)
|
|
src_shape = self.shape_of(src)
|
|
if dest_shape != src_shape:
|
|
src = self.block_builder.emit(relax.op.broadcast_to(src, dest_shape))
|
|
|
|
self.env[node.args[0]] = src
|
|
return src
|
|
|
|
def _masked_scatter(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
mask = self.env[node.args[1]]
|
|
source = self.env[node.args[2]]
|
|
ndim = len(mask.ty.shape)
|
|
if ndim == 1:
|
|
index = self.block_builder.emit(relax.op.cumsum(mask, 0, dtype="int32"))
|
|
index = self.block_builder.emit(relax.op.subtract(index, relax.const(1, "int32")))
|
|
gathered_source = self.block_builder.emit(relax.op.take(source, index, axis=0))
|
|
else:
|
|
f_mask = self.block_builder.emit(relax.op.reshape(mask, [-1]))
|
|
index = self.block_builder.emit(relax.op.cumsum(f_mask, 0, dtype="int32"))
|
|
index = self.block_builder.emit(relax.op.subtract(index, relax.const(1, "int32")))
|
|
source_shape = [-1] + [s for idx, s in enumerate(source.ty.shape) if idx >= ndim]
|
|
f_source = self.block_builder.emit(relax.op.reshape(source, source_shape))
|
|
gathered_source = self.block_builder.emit(relax.op.take(f_source, index, axis=0))
|
|
gathered_source = self.block_builder.emit(relax.op.reshape(gathered_source, x.ty.shape))
|
|
if ndim != len(x.ty.shape):
|
|
mask = self.block_builder.emit(relax.op.broadcast_to(mask, x.ty.shape))
|
|
return self.block_builder.emit(relax.op.where(mask, gathered_source, x))
|
|
|
|
def _one_hot(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
num_classes = node.args[1] if len(node.args) > 1 else node.kwargs.get("num_classes")
|
|
if num_classes is None:
|
|
raise ValueError("num_classes not found in node.args or node.kwargs")
|
|
on_value = node.args[2] if len(node.args) > 2 else node.kwargs.get("on_value", 1)
|
|
off_value = node.args[3] if len(node.args) > 3 else node.kwargs.get("off_value", 0)
|
|
axis = node.args[4] if len(node.args) > 4 else node.kwargs.get("axis", -1)
|
|
on_value = relax.prim_value(on_value)
|
|
off_value = relax.prim_value(off_value)
|
|
return self.block_builder.emit(relax.op.one_hot(x, on_value, off_value, num_classes, axis))
|
|
|
|
def _tensor(self, node: fx.Node) -> relax.Var:
|
|
dtype = node.kwargs.get("dtype", None)
|
|
if isinstance(node.args[0], float):
|
|
return relax.const(node.args[0], dtype if dtype is not None else "float32")
|
|
elif isinstance(node.args[0], int):
|
|
return relax.const(node.args[0], dtype if dtype is not None else "int64")
|
|
raise ValueError("torch.tensor with value not a float or int is not accepted")
|
|
|
|
########## DataType ##########
|
|
|
|
def _float(self, node: fx.Node) -> relax.Var:
|
|
return self.block_builder.emit(relax.op.astype(self.env[node.args[0]], "float32"))
|
|
|
|
def _half(self, node: fx.Node) -> relax.Var:
|
|
return self.block_builder.emit(relax.op.astype(self.env[node.args[0]], "float16"))
|
|
|
|
def _is_floating_point(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
return relax.const(
|
|
x.ty.dtype.dtype in ["float16", "float32", "float64", "bfloat16"], "bool"
|
|
)
|
|
|
|
def _type(self, node: fx.Node) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
dtype = TorchFXImporter._convert_data_type(node.args[1], self.env)
|
|
return self.block_builder.emit(relax.op.astype(x, dtype))
|
|
|
|
########## Others ##########
|
|
|
|
def _getattr(self, node: fx.Node) -> relax.Var:
|
|
if isinstance(self.env[node.args[0]], relax.Expr):
|
|
if node.args[1] == "dtype":
|
|
return self.env[node.args[0]].ty.dtype.dtype
|
|
elif node.args[1] == "shape":
|
|
return self.shape_of(self.env[node.args[0]])
|
|
return getattr(self.env[node.args[0]], node.args[1])
|
|
|
|
def create_input_vars(self, input_info: list[tuple[tuple[int], str]]) -> list[relax.Var]:
|
|
inputs = list()
|
|
for idx, (shape, dtype) in enumerate(input_info):
|
|
inputs.append(
|
|
relax.Var(f"inp_{idx}", relax.TensorType(shape, self._convert_data_type(dtype)))
|
|
)
|
|
return inputs
|
|
|
|
def create_convert_map(
|
|
self,
|
|
) -> dict[torch.nn.Module | str, Callable[[fx.Node], relax.Var]]:
|
|
import operator
|
|
|
|
import torch # type: ignore
|
|
from torch import nn
|
|
|
|
return {
|
|
## call_module
|
|
# unary
|
|
nn.CELU: self._celu,
|
|
nn.Dropout: lambda node: self.env[node.args[0]],
|
|
nn.ELU: self._elu,
|
|
nn.GELU: self._gelu,
|
|
nn.Hardsigmoid: self._hardsigmoid,
|
|
nn.Hardswish: self._hardswish,
|
|
nn.Hardtanh: self._hardtanh,
|
|
nn.Identity: lambda node: self.env[node.args[0]],
|
|
nn.LeakyReLU: self._leakyrelu_module,
|
|
nn.LogSoftmax: self._log_softmax_module,
|
|
nn.PReLU: self._prelu_module,
|
|
nn.ReLU: self._unary_op(relax.op.nn.relu),
|
|
nn.ReLU6: self._unary_op(relax.op.nn.relu6),
|
|
nn.Sigmoid: self._unary_op(relax.op.sigmoid),
|
|
nn.SELU: self._unary_op(relax.op.nn.selu),
|
|
nn.SiLU: self._unary_op(relax.op.nn.silu),
|
|
nn.Softmax: self._softmax_module,
|
|
nn.Softplus: self._softplus_module,
|
|
nn.Tanh: self._unary_op(relax.op.tanh),
|
|
# neural network
|
|
nn.AdaptiveAvgPool1d: self._adaptive_avg_pool1d_module,
|
|
nn.AdaptiveAvgPool2d: self._adaptive_avg_pool2d_module,
|
|
nn.AdaptiveAvgPool3d: self._adaptive_avg_pool3d_module,
|
|
nn.AvgPool1d: self._avg_pool1d_module,
|
|
nn.AvgPool2d: self._avg_pool2d_module,
|
|
nn.AvgPool3d: self._avg_pool3d_module,
|
|
nn.BatchNorm2d: self._batch_norm_2d_module,
|
|
nn.InstanceNorm1d: self._instance_norm,
|
|
nn.InstanceNorm2d: self._instance_norm,
|
|
nn.InstanceNorm3d: self._instance_norm,
|
|
nn.Conv1d: self._conv1d_module,
|
|
nn.Conv2d: self._conv2d_module,
|
|
nn.Conv3d: self._conv3d_module,
|
|
nn.ConvTranspose1d: self._conv_transpose1d_module,
|
|
nn.ConvTranspose2d: self._conv_transpose2d_module,
|
|
nn.CrossEntropyLoss: self._cross_entropy_module,
|
|
nn.GroupNorm: self._group_norm_module,
|
|
nn.LayerNorm: self._layer_norm_module,
|
|
nn.Linear: self._linear_module,
|
|
nn.MaxPool1d: self._max_pool1d_module,
|
|
nn.MaxPool2d: self._max_pool2d_module,
|
|
nn.MaxPool3d: self._max_pool3d_module,
|
|
nn.modules.sparse.Embedding: self._embedding_module,
|
|
nn.PixelShuffle: self._pixel_shuffle_module,
|
|
# tensor manipulation
|
|
nn.Flatten: self._flatten_module,
|
|
## call_function and call_method
|
|
# unary
|
|
"abs": self._unary_op(relax.op.abs),
|
|
"acos": self._unary_op(relax.op.acos),
|
|
"acosh": self._unary_op(relax.op.acosh),
|
|
"asin": self._unary_op(relax.op.asin),
|
|
"asinh": self._unary_op(relax.op.asinh),
|
|
"atan": self._unary_op(relax.op.atan),
|
|
"atanh": self._unary_op(relax.op.atanh),
|
|
"bitwise_not": self._unary_op(relax.op.bitwise_not),
|
|
"ceil": self._unary_op(relax.op.ceil),
|
|
"celu": self._celu,
|
|
"clamp": self._clamp,
|
|
"cos": self._unary_op(relax.op.cos),
|
|
"cosh": self._unary_op(relax.op.cosh),
|
|
"dropout": lambda node: self.env[node.args[0]],
|
|
"elu": self._elu,
|
|
"erf": self._unary_op(relax.op.erf),
|
|
"exp": self._unary_op(relax.op.exp),
|
|
"floor": self._unary_op(relax.op.floor),
|
|
"gelu": self._gelu,
|
|
"hardsigmoid": self._hardsigmoid,
|
|
"hardswish": self._hardswish,
|
|
"hardtanh": self._hardtanh,
|
|
"isfinite": self._unary_op(relax.op.isfinite),
|
|
"isinf": self._unary_op(relax.op.isinf),
|
|
"isin": self._isin,
|
|
"isnan": self._unary_op(relax.op.isnan),
|
|
"leaky_relu": self._leakyrelu,
|
|
"log": self._unary_op(relax.op.log),
|
|
"log2": self._log2,
|
|
"log10": self._log10,
|
|
"log1p": self._log1p,
|
|
"logical_and": self._logical_and,
|
|
"logical_not": self._logical_not,
|
|
"logical_or": self._logical_or,
|
|
"logical_xor": self._logical_xor,
|
|
"log_softmax": self._log_softmax,
|
|
"neg": self._unary_op(relax.op.negative),
|
|
"pad": self._pad,
|
|
"pixel_shuffle": self._pixel_shuffle,
|
|
"prelu": self._prelu,
|
|
"reciprocal": self._reciprocal,
|
|
"relu": self._unary_op(relax.op.nn.relu),
|
|
"relu6": self._unary_op(relax.op.nn.relu6),
|
|
"round": self._round,
|
|
"rsqrt": self._rsqrt,
|
|
"selu": self._unary_op(relax.op.nn.selu),
|
|
"sigmoid": self._unary_op(relax.op.sigmoid),
|
|
"sign": self._unary_op(relax.op.sign),
|
|
"silu": self._unary_op(relax.op.nn.silu),
|
|
"sin": self._unary_op(relax.op.sin),
|
|
"sinh": self._unary_op(relax.op.sinh),
|
|
"softmax": self._softmax,
|
|
"softplus": self._softplus,
|
|
"sqrt": self._sqrt,
|
|
"square": self._unary_op(relax.op.square),
|
|
"tan": self._unary_op(relax.op.tan),
|
|
"tanh": self._unary_op(relax.op.tanh),
|
|
"tril_": self._inplace_tril_triu(relax.op.tril),
|
|
"tril": self._tril_triu(relax.op.tril),
|
|
"triu_": self._inplace_tril_triu(relax.op.triu),
|
|
"triu": self._tril_triu(relax.op.triu),
|
|
"trunc": self._unary_op(relax.op.trunc),
|
|
# binary
|
|
"add": self._binary_op(relax.op.add, operator.add),
|
|
"and_": self._binary_op(relax.op.bitwise_and, operator.and_),
|
|
"atan2": self._binary_op(relax.op.atan2, torch.atan2),
|
|
"bitwise_or_": self._binary_op_inplace(relax.op.bitwise_or, operator.or_),
|
|
"bitwise_or": self._binary_op(relax.op.bitwise_or, operator.or_),
|
|
"div": self._div,
|
|
"eq": self._binary_op(relax.op.equal, operator.eq),
|
|
"floordiv": self._binary_op(relax.op.floor_divide, operator.floordiv),
|
|
"fmod": self._fmod,
|
|
"ge": self._binary_op(relax.op.greater_equal, operator.ge),
|
|
"gt": self._binary_op(relax.op.greater, operator.gt),
|
|
"iadd": self._binary_op(relax.op.add, operator.add),
|
|
"le": self._binary_op(relax.op.less_equal, operator.le),
|
|
"lshift": self._binary_op(relax.op.left_shift, operator.lshift),
|
|
"lt": self._binary_op(relax.op.less, operator.lt),
|
|
"matmul": self._binary_op(
|
|
partial(relax.op.linear_algebra.matmul, out_dtype="float32"), operator.matmul
|
|
),
|
|
"max": self._binary_op(relax.op.maximum, max),
|
|
"min": self._binary_op(relax.op.minimum, min),
|
|
"mod": self._binary_op(relax.op.floor_mod, operator.mod),
|
|
"mul": self._binary_op(relax.op.multiply, operator.mul),
|
|
"ne": self._binary_op(relax.op.not_equal, operator.ne),
|
|
"outer": lambda node: self.block_builder.emit(
|
|
relax.op.outer(self.env[node.args[0]], self.env[node.args[1]])
|
|
),
|
|
"pow": self._pow,
|
|
"or_": self._binary_op(relax.op.bitwise_or, operator.or_),
|
|
"rshift": self._binary_op(relax.op.right_shift, operator.rshift),
|
|
"rsub": self._rsub,
|
|
"sub": self._binary_op(relax.op.subtract, operator.sub),
|
|
"truediv": self._binary_op(relax.op.divide, operator.truediv),
|
|
"xor": self._binary_op(relax.op.bitwise_xor, operator.xor),
|
|
# neural network
|
|
"adaptive_avg_pool1d": self._adaptive_avg_pool1d,
|
|
"adaptive_avg_pool2d": self._adaptive_avg_pool2d,
|
|
"adaptive_avg_pool3d": self._adaptive_avg_pool3d,
|
|
"addmm": self._addmm,
|
|
"avg_pool1d": self._avg_pool1d,
|
|
"avg_pool2d": self._avg_pool2d,
|
|
"avg_pool3d": self._avg_pool3d,
|
|
"baddbmm": self._baddbmm,
|
|
"bmm": self._binary_op(
|
|
partial(relax.op.linear_algebra.matmul, out_dtype="float32"), operator.matmul
|
|
),
|
|
"conv_transpose1d": self._conv_transpose1d,
|
|
"conv_transpose2d": self._conv_transpose2d,
|
|
"conv1d": self._conv1d,
|
|
"conv2d": self._conv2d,
|
|
"conv3d": self._conv3d,
|
|
"cross_entropy": self._cross_entropy,
|
|
"einsum": self._einsum,
|
|
"interpolate": self._interpolate,
|
|
"layer_norm": self._layer_norm,
|
|
"linear": self._linear,
|
|
"max_pool1d": self._max_pool1d,
|
|
"max_pool2d": self._max_pool2d,
|
|
"max_pool3d": self._max_pool3d,
|
|
"scaled_dot_product_attention": self._scaled_dot_product_attention,
|
|
"stochastic_depth": lambda node: self.env[node.args[0]],
|
|
"unbind": self._unbind,
|
|
# linear interpolation
|
|
"lerp": self._lerp,
|
|
# statistical
|
|
"mean": self._mean,
|
|
"norm": self._norm,
|
|
"prod": self._prod,
|
|
"std": self._std,
|
|
"sum": self._sum,
|
|
"var": self._var,
|
|
# search
|
|
"argmax": self._argmax_argmin(relax.op.argmax),
|
|
"argmin": self._argmax_argmin(relax.op.argmin),
|
|
"where": self._where,
|
|
"bucketize": self._bucketize,
|
|
# tensor manipulation
|
|
"argsort": self._argsort,
|
|
"broadcast_to": self._broadcast_to,
|
|
"cat": self._cat,
|
|
"chunk": self._chunk,
|
|
"concat": self._cat,
|
|
"contiguous": lambda node: self.env[node.args[0]],
|
|
"cumprod": self._cumprod,
|
|
"cumsum": self._cumsum,
|
|
"expand": self._expand,
|
|
"expand_as.default": self._expand_as,
|
|
"flatten": self._flatten,
|
|
"flip": self._flip,
|
|
"gather": self._gather,
|
|
"index_put_": self._index_put,
|
|
"meshgrid": self._meshgrid,
|
|
"narrow": self._narrow,
|
|
"numel": self._numel,
|
|
"permute": self._permute,
|
|
"repeat": self._repeat,
|
|
"roll": self._roll,
|
|
"reshape": self._reshape,
|
|
"scatter": self._scatter,
|
|
"select": self._select,
|
|
"size": self._size,
|
|
"slice_scatter": self._slice_scatter,
|
|
"sort": self._sort,
|
|
"split": self._split,
|
|
"squeeze": self._squeeze,
|
|
"stack": self._stack,
|
|
"take": self._take,
|
|
"tile": self._tile,
|
|
"topk": self._topk,
|
|
"transpose": self._transpose,
|
|
"unsqueeze": lambda node: self.block_builder.emit(
|
|
relax.op.expand_dims(self.env[node.args[0]], node.args[1])
|
|
),
|
|
"view": self._reshape,
|
|
# tensor creation
|
|
"arange": self._arange,
|
|
"clone": lambda node: self.env[node.args[0]],
|
|
"empty": self._empty,
|
|
"empty_like": self._empty_like,
|
|
"eye": self._eye,
|
|
"fill": self._fill,
|
|
"fill_": self._inplace_fill,
|
|
"full": self._full,
|
|
"index_select": self._index_select,
|
|
"linspace": self._linspace,
|
|
"masked_fill_": self._inplace_masked_fill,
|
|
"masked_fill": self._masked_fill,
|
|
"masked_scatter": self._masked_scatter,
|
|
"new_ones": self._new_ones,
|
|
"new_zeros": self._new_zeros,
|
|
"ones": self._ones,
|
|
"one_hot": self._one_hot,
|
|
"ones_like": lambda node: self.block_builder.emit(
|
|
relax.op.ones_like(self.env[node.args[0]])
|
|
),
|
|
"tensor": self._tensor,
|
|
"zero_": self._zeros_inplace,
|
|
"zeros_like": self._zeros_like,
|
|
"copy_": self._inplace_copy,
|
|
# datatype
|
|
"astype": self._type,
|
|
"float": self._float,
|
|
"half": self._half,
|
|
"is_floating_point": self._is_floating_point,
|
|
"to": self._to,
|
|
"type": self._type,
|
|
"type_as": self._type_as,
|
|
# other
|
|
"getattr": self._getattr,
|
|
"getitem": self._getitem,
|
|
"sym_size.int": self._sym_size_int,
|
|
"item": self._item,
|
|
}
|
|
|
|
def from_fx(
|
|
self,
|
|
model,
|
|
input_info: list[tuple[tuple[int], str]],
|
|
keep_params_as_input: bool,
|
|
unwrap_unit_return_tuple: bool,
|
|
no_bind_return_tuple: bool,
|
|
custom_convert_map: dict | None = None,
|
|
) -> tvm.IRModule:
|
|
"""Convert a PyTorch FX GraphModule to a Relax program."""
|
|
from torch import fx
|
|
|
|
if custom_convert_map:
|
|
custom_ops = set(custom_convert_map.keys())
|
|
self.update_convert_map(custom_convert_map)
|
|
else:
|
|
custom_ops = set()
|
|
self.named_modules = dict(model.named_modules())
|
|
|
|
graph: fx.Graph = model.graph
|
|
|
|
# Create input variables.
|
|
inputs = self.create_input_vars(input_info)
|
|
|
|
# Initialize the block builder with a function and a dataflow block.
|
|
func_name = "main"
|
|
self.block_builder = relax.BlockBuilder()
|
|
params = []
|
|
if keep_params_as_input:
|
|
func_attrs = {"num_input": len(inputs)}
|
|
for name, param in sorted(model.named_parameters(), key=lambda x: x[0]):
|
|
shape = param.data.shape
|
|
dtype = self._convert_data_type(str(param.data.dtype))
|
|
inputs.append(relax.Var(name, relax.TensorType(shape, dtype)))
|
|
self.params[param] = inputs[-1]
|
|
params.append(tvm.runtime.tensor(param.data.cpu().numpy()))
|
|
else:
|
|
func_attrs = None
|
|
|
|
# Find all the missing function types
|
|
self._check_unsupported_func_type(graph.nodes)
|
|
|
|
from tvm import tirx
|
|
|
|
sym_vars = {v.name: v for shape, _ in input_info for v in shape if isinstance(v, tirx.Var)}
|
|
|
|
with self.block_builder.function(name=func_name, params=inputs.copy(), attrs=func_attrs):
|
|
output = None
|
|
with self.block_builder.dataflow():
|
|
# Translate model parameters.
|
|
for _, param in model.named_parameters():
|
|
shape = param.data.shape
|
|
dtype = self._convert_data_type(str(param.data.dtype))
|
|
if dtype in ("float32", "float16"):
|
|
if not keep_params_as_input:
|
|
self.params[param] = self._convert_torch_tensor_to_relax(param)
|
|
else:
|
|
raise ValueError(f"Unsupported data type for model parameters: {dtype}")
|
|
# Translate the model.
|
|
for node in graph.nodes:
|
|
if node.op == "placeholder":
|
|
if "grapharg" in node.meta and node.meta["grapharg"].fake_tensor is None:
|
|
# Sym input: bind to the matching shape var if referenced
|
|
if node.name in sym_vars:
|
|
self.env[node] = sym_vars[node.name]
|
|
continue
|
|
|
|
assert len(inputs) > 0, "Provided inputs is less than actual inputs"
|
|
self.env[node] = inputs.pop(0)
|
|
elif node.op == "output":
|
|
args = self.retrieve_args(node)
|
|
assert len(args) == 1
|
|
|
|
# return tuple
|
|
if isinstance(args[0], tuple | list | relax.Tuple):
|
|
# unit tuple
|
|
if unwrap_unit_return_tuple and len(args[0]) == 1:
|
|
output = self.block_builder.emit_output(args[0][0])
|
|
elif no_bind_return_tuple:
|
|
output = []
|
|
for ret in args[0]:
|
|
output.append(self.block_builder.emit_output(ret))
|
|
|
|
if output is None:
|
|
output = self.block_builder.emit_output(args[0])
|
|
break
|
|
elif node.op == "get_attr":
|
|
self.env[node] = self._fetch_attr(model, node.target)
|
|
elif node.op == "call_module":
|
|
module = self.named_modules[node.target]
|
|
assert type(module) in self.convert_map, (
|
|
f"Unsupported module type {type(module)}"
|
|
)
|
|
self.env[node] = self.convert_map[type(module)](node)
|
|
elif node.op == "call_function":
|
|
func_name = node.target.__name__
|
|
if func_name in custom_ops:
|
|
self.env[node] = self.convert_map[func_name](node, self)
|
|
else:
|
|
self.env[node] = self.convert_map[func_name](node)
|
|
elif node.op == "call_method":
|
|
assert node.target in self.convert_map, (
|
|
f"Unsupported function target {node.target}"
|
|
)
|
|
self.env[node] = self.convert_map[node.target](node)
|
|
else:
|
|
raise ValueError(f"Unsupported op {node.op}")
|
|
assert output is not None
|
|
self.block_builder.emit_func_output(output)
|
|
|
|
mod = self.block_builder.get()
|
|
if keep_params_as_input:
|
|
mod["main"] = mod["main"].with_attr("params", params)
|
|
return mod
|
|
|
|
|
|
def from_fx(
|
|
model,
|
|
input_info: list[tuple[tuple[int], str]],
|
|
*,
|
|
keep_params_as_input: bool = False,
|
|
unwrap_unit_return_tuple: bool = False,
|
|
no_bind_return_tuple: bool = False,
|
|
custom_convert_map: dict | None = None,
|
|
default_image_layout: str = "NCHW",
|
|
) -> tvm.IRModule:
|
|
"""Convert a PyTorch FX GraphModule to a Relax program
|
|
|
|
Parameters
|
|
----------
|
|
model : fx.GraphModule
|
|
The PyTorch FX GraphModule to convert.
|
|
|
|
input_info : List[Tuple[Tuple[int], str]]
|
|
A list of shapes and data types of input tensors.
|
|
|
|
keep_params_as_input : bool
|
|
Whether to keep model parameters as input variables.
|
|
|
|
unwrap_unit_return_tuple : bool
|
|
A boolean flag indicating if to the return value when it is an unit tuple.
|
|
When the return value is not a unit tuple, no unwrap will take place.
|
|
|
|
no_bind_return_tuple : bool
|
|
A boolean flag indicating whether to bind the return tuple as a relax var.
|
|
If the flag is true and the return value is a tuple, it will not bind it to a var.
|
|
|
|
custom_convert_map : Dictionary of str to Relax op
|
|
A custom op conversion map in the same format as TorchFXImporter.convert_map
|
|
|
|
default_image_layout : str
|
|
The default layout for image operations (e.g., "NCHW" or "NHWC").
|
|
Default is "NCHW" which is the standard PyTorch layout.
|
|
|
|
Returns
|
|
-------
|
|
output : tvm.IRModule
|
|
The import result IRModule, with the function "main" containing the
|
|
translated logic.
|
|
If `keep_params_as_input` is true, the "main" function have an attribute
|
|
"params" that contains the weights of the input model. The weights
|
|
can be detached by `relax.frontend.detach_params`.
|
|
|
|
Examples
|
|
--------
|
|
Users can use the FX tracer or dynamo.export() to extract
|
|
a fx.GraphModule from a PyTorch model. The following codes show
|
|
how to convert a PyTorch model to a Relax program.
|
|
|
|
.. code-block:: python
|
|
|
|
# Import the importer.
|
|
import numpy as np
|
|
import torch
|
|
from tvm.relax.frontend.torch_fx import from_fx
|
|
from torch import _dynamo as dynamo
|
|
|
|
# Define the module
|
|
class MyModule(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = torch.nn.Linear(in_features=10, out_features=7, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.linear(input)
|
|
|
|
# Instantiate the model and create the input info dict.
|
|
torch_model = MyModule()
|
|
input_info = [((128, 10), "float32")]
|
|
input_tensors = [
|
|
torch.astensor(np.random.randn(*shape).astype(dtype))
|
|
for shape, dtype in input_info
|
|
]
|
|
|
|
# Use FX tracer to trace the PyTorch model.
|
|
graph_module = fx.symbolic_trace(torch_model)
|
|
|
|
# Use the dynamo.export() to export the PyTorch model to FX.
|
|
try:
|
|
graph_module = dynamo.export(torch_model, *input_tensors)
|
|
except Exception:
|
|
raise RuntimeError("Failed to export the PyTorch model to FX.")
|
|
|
|
# Use the importer to import the PyTorch model to Relax.
|
|
mod: tvm.IRModule = from_fx(graph_module, input_info)
|
|
|
|
# Print out the imported model.
|
|
print(mod.script())
|
|
|
|
Notes
|
|
-----
|
|
For a given PyTorch model, to lookup the names of the model inputs in
|
|
FX, one can use
|
|
|
|
.. code-block:: python
|
|
|
|
fx.symbolic_trace(model).graph.print_tabular()
|
|
|
|
to print out the tabular representation of the PyTorch module, and then
|
|
check the placeholder rows in the beginning of the tabular.
|
|
"""
|
|
return TorchFXImporter(default_image_layout=default_image_layout).from_fx(
|
|
model,
|
|
input_info,
|
|
keep_params_as_input,
|
|
unwrap_unit_return_tuple,
|
|
no_bind_return_tuple,
|
|
custom_convert_map=custom_convert_map,
|
|
)
|