[GitHub] [incubator-tvm] zhiics commented on a change in pull request #6578: [tvmc] Introduce 'run' subcommand (part 4/4)

[GitBox <gi...@apache.org>]

zhiics commented on a change in pull request #6578:
URL: https://github.com/apache/incubator-tvm/pull/6578#discussion_r496142487



##########
File path: python/tvm/driver/tvmc/runner.py
##########
@@ -0,0 +1,450 @@
+# 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.
+"""
+Provides support to run compiled networks both locally and remotely.
+"""
+import json
+import logging
+import os
+import tarfile
+import tempfile
+
+import numpy as np
+import tvm
+from tvm import rpc
+from tvm.autotvm.measure import request_remote
+from tvm.contrib import graph_runtime as runtime
+from tvm.contrib.debugger import debug_runtime
+
+from . import common
+from .common import TVMCException
+from .main import register_parser
+
+
+# pylint: disable=invalid-name
+logger = logging.getLogger("TVMC")
+
+
+@register_parser
+def add_run_parser(subparsers):
+    """ Include parser for 'run' subcommand """
+
+    parser = subparsers.add_parser("run", help="run a compiled module")
+    parser.set_defaults(func=drive_run)
+
+    # TODO --device needs to be extended and tested to support other targets,
+    #      like 'cl', 'webgpu', etc (@leandron)
+    parser.add_argument(
+        "--device",
+        choices=["cpu", "gpu"],
+        default="cpu",
+        help="target device to run the compiled module",
+    )
+    parser.add_argument(
+        "--fill-mode",
+        choices=["zeros", "ones", "random"],
+        default="zeros",
+        help="fill all input tensors with values",
+    )
+    parser.add_argument("-i", "--inputs", help="path to the .npz input file")
+    parser.add_argument("-o", "--outputs", help="path to the .npz output file")
+    parser.add_argument(
+        "--print-time", action="store_true", help="record and print the execution time(s)"
+    )
+    parser.add_argument(
+        "--print-top",
+        metavar="N",
+        type=int,
+        help="print the top n values and indices of the output tensor",
+    )
+    parser.add_argument(
+        "--profile", action="store_true", help="generate profiling data from the runtime execution"
+    )
+    parser.add_argument("--repeat", metavar="N", type=int, default=1, help="repeat the run n times")
+    parser.add_argument(
+        "--rpc-key",
+        nargs=1,
+        help="the RPC tracker key of the target device",
+    )
+    parser.add_argument(
+        "--rpc-tracker",
+        nargs=1,
+        help="hostname (required) and port (optional, defaults to 9090) of the RPC tracker, "
+        "e.g. '192.168.0.100:9999'",
+    )
+    parser.add_argument("FILE", help="path to the compiled module file")
+
+
+def drive_run(args):
+    """Invoke runner module with command line arguments
+
+    Parameters
+    ----------
+    args: argparse.Namespace
+        Arguments from command line parser.
+    """
+    inputs = {}
+    if args.inputs:
+        inputs = np.load(args.inputs)
+
+    rpc_hostname, rpc_port = common.tracker_host_port_from_cli(args.rpc_tracker)
+
+    outputs, times = run_module(
+        args.FILE,
+        rpc_hostname,
+        rpc_port,
+        args.rpc_key,
+        inputs=inputs,
+        device=args.device,
+        fill_mode=args.fill_mode,
+        repeat=args.repeat,
+        profile=args.profile,
+    )
+
+    if args.print_time:
+        stat_table = format_times(times)
+        # print here is intentional
+        print(stat_table)
+
+    if args.print_top:
+        top_results = get_top_results(outputs, args.print_top)
+        # print here is intentional
+        print(top_results)
+
+    if args.outputs:
+        # Save the outputs
+        np.savez(args.outputs, **outputs)
+
+
+def get_input_info(graph_str, params):
+    """Return the 'shape' and 'dtype' dictionaries for the input
+    tensors of a compiled module.
+
+    Parameters
+    ----------
+    graph_str : str
+        JSON graph of the module serialized as a string.
+    params : bytearray
+        Params serialized as a bytearray.
+
+    Returns
+    -------
+    shape_dict : dict
+        Shape dictionary - {input_name: tuple}.
+    dtype_dict : dict
+        dtype dictionary - {input_name: dtype}.
+    """
+    # NOTE - We can't simply get the input tensors from a TVM graph
+    # because weight tensors are treated equivalently. Therefore, to
+    # find the input tensors we look at the 'arg_nodes' in the graph
+    # (which are either weights or inputs) and check which ones don't
+    # appear in the params (where the weights are stored). These nodes
+    # are therefore inferred to be input tensors.

Review comment:
       This logic is okay to me for now. In the long run, we may want to change graph runtime a bit and let it save inputs and params separately. params can be obtained from MetadataModule directly.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org