PyTorch Relay 推理

目录

PyTorch Relay 推理#

import sys

sys.path.extend(["../../../../tests"])
# from tools.tag_span import _create_span, _set_span, _verify_structural_equal_with_span
from tools.torch_utils import verify_model

import logging
import numpy as np
import torch
import tvm
logging.basicConfig(level=logging.DEBUG)

def list_ops(expr):
    """list_ops"""

    class OpLister(tvm.relay.ExprVisitor):
        """OpLister inherits from ExprVisitor"""

        def visit_op(self, op):
            if op not in self.node_set:
                self.node_list.append(op)
            return super().visit_op(op)

        def list_nodes(self, expr):
            self.node_set = {}
            self.node_list = []
            self.visit(expr)
            return self.node_list

    return OpLister().list_nodes(expr)

torch.matmul#

class MatMul1(torch.nn.Module):
    def forward(self, *args):
        return torch.matmul(args[0], args[1])

# vector x vector - 1D x 1D
tensor1 = torch.randn(4)
tensor2 = torch.randn(4)
verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.dense"])

DEBUG:autotvm:Finish loading 35 records
WARNING:autotvm:One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.
DEBUG:autotvm:Cannot find tuning records for:
    target=llvm -keys=cpu -mtriple=x86_64-unknown-linux-gnu
    key=('dense_pack.x86', ('TENSOR', (1, 4), 'float32'), ('TENSOR', (1, 4), 'float32'), None, 'float32')
TVM will apply a default schedule which may negatively impact performance.
INFO:te_compiler:Using dense_pack.x86 for nn.dense based on highest priority (10)
INFO:te_compiler:Using injective.cpu for expand_dims based on highest priority (10)
INFO:te_compiler:Using injective.cpu for expand_dims based on highest priority (10)
INFO:te_compiler:Using transpose.generic for transpose based on highest priority (10)
INFO:te_compiler:Using layout_transform.generic for layout_transform based on highest priority (10)
DEBUG:autotvm:Cannot find tuning records for:
    target=llvm -keys=cpu -mtriple=x86_64-unknown-linux-gnu
    key=('dense_pack.x86', ('TENSOR', (1, 4), 'float32'), ('TENSOR', (1, 4, 1), 'float32'), None, 'float32')
TVM will apply a default schedule which may negatively impact performance.
INFO:te_compiler:Using dense_pack.x86 for nn.contrib_dense_pack based on highest priority (10)
INFO:te_compiler:Using injective.cpu for squeeze based on highest priority (10)
目录