How to set parameter into pipeline module

popojames · November 2, 2021, 4:44am

Hello @hijiang,

I have followed your work and made some extensions to support such splitting as I mentioned in the previous discussions Setting the CPU affinity and number of cores locally without RPC Remote - #4 by popojames. Can TVM split work into different layers, and assign layers into different cores? - #10 by hjiang. I was able to split the network and with my own CPU affinity setting.

Now I try to assign param of the pre-trained model into the pipeline module, I noticed the function of PipelineModule.set_input support such setting:

github.com

apache/tvm/blob/fba531e55f032367acc28090605c5533b2e88f55/python/tvm/contrib/pipeline_executor.py#L222


        the submodule index


    params : dict of str to NDArray
        Additional arguments
    """
    assert mod_idx >= 0
    self._set_input(key, tvm.nd.array(value, tvm.cpu()), mod_idx)


    if params:
        for param in params:
            self.graph_modules_[mod_idx].set_input(**param)


def run(self):
    """Run forward execution of the graph"""
    self._run()


def stop(self):
    """Stop pipeline run"""
    self._stop()


def get_num_outputs(self):

However, this function will induce the following error:

AttributeError: 'Module' object has no attribute 'set_input'

This is because of the type of self.graph_modules_[mod_idx] is something like Module(GraphExecutor, 27517208) instead of tvm.contrib.graph_executor.GraphModule. In this case, the latter can access set_input but the former cannot.

Is there any other way to assign parameters or can I transfer from Module(GraphExecutor, XXX) to tvm.contrib.graph_executor.GraphModule in this case to set the parameters?

Thanks for your help in advance and many thanks for developing this TVM function.

hjiang · November 2, 2021, 7:00pm

@popojames , thanks for trying the pipeline runtime, the set params issue will get address in our new patch, you also can try following change as a quick fix.

diff --git a/python/tvm/contrib/pipeline_executor.py b/python/tvm/contrib/pipeline_executor.py
index bd8c9a768..c1f7842b3 100644
--- a/python/tvm/contrib/pipeline_executor.py
+++ b/python/tvm/contrib/pipeline_executor.py
@@ -157,7 +157,7 @@ def create(pipeline_mods, mod_config):
         mod = graph_executor.GraphModule(
             pipeline_mod["default"](pipeline_mods[pipeline_mod]["dev"])
         )
-        mods.append(mod.module)
+        mods.append(mod)

     submodule = PipelineModule(mods, json.dumps(mod_config))
     # submodule = PipelineModule(pipeline_mods, json.dumps(mod_config))
@@ -181,7 +181,7 @@ class PipelineModule(object):
     def __init__(self, modules, pipeline_config):
         mods = []
         for module in modules:
-            mods.append(module)
+            mods.append(module.module)

         pipelinecreate = tvm._ffi.get_global_func("tvm.pipeline_executor.create")
         assert pipelinecreate

popojames · November 2, 2021, 9:28pm

Thanks for replying, I will try and keep an eye on the new patches.

popojames · November 3, 2021, 4:05am

Hello @hjiang,

Thanks for answering, I was able to feed the params into the model by adopting your change.

But I add one more change: I changed to the following code

if params:
     for param in params:
         self.graph_modules_[mod_idx].set_input(**params)

into the following code to avoid the error

if params:
     self.graph_modules_[mod_idx].set_input(**params)

github.com/apache/tvm

python/tvm/contrib/pipeline_executor.py

1ee5831b6


      
          
              mod_idx : int
                  the submodule index
          
              params : dict of str to NDArray
                  Additional arguments
              """
              assert mod_idx >= 0
              self._set_input(key, tvm.nd.array(value, tvm.cpu()), mod_idx)
          
              if params:
                  for param in params:
                      self.graph_modules_[mod_idx].set_input(**param)
          
          def run(self):
              """Run forward execution of the graph"""
              self._run()
          
          def stop(self):
              """Stop pipeline run"""
              self._stop()

====================================================================

More info: What I want to do is very similar to do what you have done in tvm/test_pipeline_executor.py at fba531e55f032367acc28090605c5533b2e88f55 · apache/tvm (github.com) To be more specific, I tried to compare the correctness of the pipeline setting, in another word, if I feed the input to with pipeline and run the inference, I was expected to get the same result as I get in without the pipeline setting. Now I was able to achieve so.

Thanks again