# 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, )