pydoc/2.58.0/apache_beam.ml.inference.pytorch_inference.html - beam-site - Git at Google



 <!DOCTYPE html>
 <!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
 <!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
 <head>
   <meta charset="utf-8">

   <meta name="viewport" content="width=device-width, initial-scale=1.0">

   <title>apache_beam.ml.inference.pytorch_inference module &mdash; Apache Beam 2.58.0 documentation</title>


   <script type="text/javascript" src="_static/js/modernizr.min.js"></script>


       <script type="text/javascript" id="documentation_options" data-url_root="./" src="_static/documentation_options.js"></script>
         <script type="text/javascript" src="_static/jquery.js"></script>
         <script type="text/javascript" src="_static/underscore.js"></script>
         <script type="text/javascript" src="_static/doctools.js"></script>
         <script type="text/javascript" src="_static/language_data.js"></script>
         <script async="async" type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>

     <script type="text/javascript" src="_static/js/theme.js"></script>


   <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
   <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
     <link rel="index" title="Index" href="genindex.html" />
     <link rel="search" title="Search" href="search.html" />
     <link rel="next" title="apache_beam.ml.inference.sklearn_inference module" href="apache_beam.ml.inference.sklearn_inference.html" />
     <link rel="prev" title="apache_beam.ml.inference.onnx_inference module" href="apache_beam.ml.inference.onnx_inference.html" />
 </head>

 <body class="wy-body-for-nav">


   <div class="wy-grid-for-nav">

     <nav data-toggle="wy-nav-shift" class="wy-nav-side">
       <div class="wy-side-scroll">
         <div class="wy-side-nav-search" >


             <a href="index.html" class="icon icon-home"> Apache Beam


           </a>


               <div class="version">
                 2.58.0
               </div>


 <div role="search">
   <form id="rtd-search-form" class="wy-form" action="search.html" method="get">
     <input type="text" name="q" placeholder="Search docs" />
     <input type="hidden" name="check_keywords" value="yes" />
     <input type="hidden" name="area" value="default" />
   </form>
 </div>


         </div>

         <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">


               <ul class="current">
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.coders.html">apache_beam.coders package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.dataframe.html">apache_beam.dataframe package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.io.html">apache_beam.io package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.metrics.html">apache_beam.metrics package</a></li>
 <li class="toctree-l1 current"><a class="reference internal" href="apache_beam.ml.html">apache_beam.ml package</a><ul class="current">
 <li class="toctree-l2 current"><a class="reference internal" href="apache_beam.ml.html#subpackages">Subpackages</a><ul class="current">
 <li class="toctree-l3"><a class="reference internal" href="apache_beam.ml.gcp.html">apache_beam.ml.gcp package</a></li>
 <li class="toctree-l3 current"><a class="reference internal" href="apache_beam.ml.inference.html">apache_beam.ml.inference package</a><ul class="current">
 <li class="toctree-l4 current"><a class="reference internal" href="apache_beam.ml.inference.html#submodules">Submodules</a></li>
 </ul>
 </li>
 <li class="toctree-l3"><a class="reference internal" href="apache_beam.ml.transforms.html">apache_beam.ml.transforms package</a></li>
 </ul>
 </li>
 </ul>
 </li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.options.html">apache_beam.options package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.portability.html">apache_beam.portability package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.runners.html">apache_beam.runners package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.testing.html">apache_beam.testing package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.transforms.html">apache_beam.transforms package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.typehints.html">apache_beam.typehints package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.utils.html">apache_beam.utils package</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.yaml.html">apache_beam.yaml package</a></li>
 </ul>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.error.html">apache_beam.error module</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.pipeline.html">apache_beam.pipeline module</a></li>
 <li class="toctree-l1"><a class="reference internal" href="apache_beam.pvalue.html">apache_beam.pvalue module</a></li>
 </ul>


         </div>
       </div>
     </nav>

     <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">


       <nav class="wy-nav-top" aria-label="top navigation">

           <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
           <a href="index.html">Apache Beam</a>

       </nav>


       <div class="wy-nav-content">

         <div class="rst-content">


 <div role="navigation" aria-label="breadcrumbs navigation">

   <ul class="wy-breadcrumbs">

       <li><a href="index.html">Docs</a> &raquo;</li>

           <li><a href="apache_beam.ml.html">apache_beam.ml package</a> &raquo;</li>

           <li><a href="apache_beam.ml.inference.html">apache_beam.ml.inference package</a> &raquo;</li>

       <li>apache_beam.ml.inference.pytorch_inference module</li>


       <li class="wy-breadcrumbs-aside">


             <a href="_sources/apache_beam.ml.inference.pytorch_inference.rst.txt" rel="nofollow"> View page source</a>


       </li>

   </ul>


   <hr/>
 </div>
           <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
            <div itemprop="articleBody">

   <div class="section" id="module-apache_beam.ml.inference.pytorch_inference">
 <span id="apache-beam-ml-inference-pytorch-inference-module"></span><h1>apache_beam.ml.inference.pytorch_inference module<a class="headerlink" href="#module-apache_beam.ml.inference.pytorch_inference" title="Permalink to this headline">¶</a></h1>
 <dl class="class">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor">
 <em class="property">class </em><code class="descclassname">apache_beam.ml.inference.pytorch_inference.</code><code class="descname">PytorchModelHandlerTensor</code><span class="sig-paren">(</span><em>state_dict_path: Optional[str] = None, model_class: Optional[Callable[[...], &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9154a0bfd0&gt;]] = None, model_params: Optional[Dict[str, Any]] = None, device: str = 'CPU', *, inference_fn: Callable[[Sequence[&lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9154a0b610&gt;], &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9154a0b700&gt;, &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9154a0bc10&gt;, Optional[Dict[str, Any]], Optional[str]], Iterable[apache_beam.ml.inference.base.PredictionResult]] = &lt;function default_tensor_inference_fn&gt;, torch_script_model_path: Optional[str] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, load_model_args: Optional[Dict[str, Any]] = None, **kwargs</em><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor" title="Permalink to this definition">¶</a></dt>
 <dd><p>Bases: <a class="reference internal" href="apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler" title="apache_beam.ml.inference.base.ModelHandler"><code class="xref py py-class docutils literal notranslate"><span class="pre">apache_beam.ml.inference.base.ModelHandler</span></code></a></p>
 <p>Implementation of the ModelHandler interface for PyTorch.</p>
 <dl class="docutils">
 <dt>Example Usage for torch model::</dt>
 <dd><dl class="first last docutils">
 <dt>pcoll | RunInference(PytorchModelHandlerTensor(state_dict_path=”my_uri”,</dt>
 <dd>model_class=”my_class”))</dd>
 </dl>
 </dd>
 <dt>Example Usage for torchscript model::</dt>
 <dd><dl class="first last docutils">
 <dt>pcoll | RunInference(PytorchModelHandlerTensor(</dt>
 <dd>torch_script_model_path=”my_uri”))</dd>
 </dl>
 </dd>
 </dl>
 <p>See <a class="reference external" href="https://pytorch.org/tutorials/beginner/saving_loading_models.html">https://pytorch.org/tutorials/beginner/saving_loading_models.html</a>
 for details</p>
 <table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Parameters:</th><td class="field-body"><ul class="first last simple">
 <li><strong>state_dict_path</strong> – path to the saved dictionary of the model state.</li>
 <li><strong>model_class</strong> – class of the Pytorch model that defines the model
 structure.</li>
 <li><strong>model_params</strong> – A dictionary of arguments required to instantiate the model
 class.</li>
 <li><strong>device</strong> – the device on which you wish to run the model. If
 <code class="docutils literal notranslate"><span class="pre">device</span> <span class="pre">=</span> <span class="pre">GPU</span></code> then a GPU device will be used if it is available.
 Otherwise, it will be CPU.</li>
 <li><strong>inference_fn</strong> – the inference function to use during RunInference.
 default=_default_tensor_inference_fn</li>
 <li><strong>torch_script_model_path</strong> – <dl class="docutils">
 <dt>Path to the torch script model.</dt>
 <dd>the model will be loaded using <cite>torch.jit.load()</cite>.</dd>
 <dt><cite>state_dict_path</cite>, <cite>model_class</cite> and <cite>model_params</cite></dt>
 <dd>arguments will be disregarded.</dd>
 </dl>
 </li>
 <li><strong>min_batch_size</strong> – the minimum batch size to use when batching inputs. This
 batch will be fed into the inference_fn as a Sequence of Tensors.</li>
 <li><strong>max_batch_size</strong> – the maximum batch size to use when batching inputs. This
 batch will be fed into the inference_fn as a Sequence of Tensors.</li>
 <li><strong>max_batch_duration_secs</strong> – the maximum amount of time to buffer a batch
 before emitting; used in streaming contexts.</li>
 <li><strong>large_model</strong> – set to true if your model is large enough to run into
 memory pressure if you load multiple copies. Given a model that
 consumes N memory and a machine with W cores and M memory, you should
 set this to True if N*W &gt; M.</li>
 <li><strong>model_copies</strong> – The exact number of models that you would like loaded
 onto your machine. This can be useful if you exactly know your CPU or
 GPU capacity and want to maximize resource utilization.</li>
 <li><strong>load_model_args</strong> – a dictionary of parameters passed to the torch.load
 function to specify custom config for loading models.</li>
 <li><strong>kwargs</strong> – ‘env_vars’ can be used to set environment variables
 before loading the model.</li>
 </ul>
 </td>
 </tr>
 </tbody>
 </table>
 <p><strong>Supported Versions:</strong> RunInference APIs in Apache Beam have been tested
 with PyTorch 1.9 and 1.10.</p>
 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.load_model">
 <code class="descname">load_model</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9157567e50&gt;<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.load_model"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.load_model" title="Permalink to this definition">¶</a></dt>
 <dd><p>Loads and initializes a Pytorch model for processing.</p>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.update_model_path">
 <code class="descname">update_model_path</code><span class="sig-paren">(</span><em>model_path: Optional[str] = None</em><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.update_model_path"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.update_model_path" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.run_inference">
 <code class="descname">run_inference</code><span class="sig-paren">(</span><em>batch: Sequence[&lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9154a0b280&gt;], model: &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9154a0be20&gt;, inference_args: Optional[Dict[str, Any]] = None</em><span class="sig-paren">)</span> &#x2192; Iterable[apache_beam.ml.inference.base.PredictionResult]<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.run_inference"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.run_inference" title="Permalink to this definition">¶</a></dt>
 <dd><p>Runs inferences on a batch of Tensors and returns an Iterable of
 Tensor Predictions.</p>
 <p>This method stacks the list of Tensors in a vectorized format to optimize
 the inference call.</p>
 <table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Parameters:</th><td class="field-body"><ul class="first simple">
 <li><strong>batch</strong> – A sequence of Tensors. These Tensors should be batchable, as this
 method will call <cite>torch.stack()</cite> and pass in batched Tensors with
 dimensions (batch_size, n_features, etc.) into the model’s forward()
 function.</li>
 <li><strong>model</strong> – A PyTorch model.</li>
 <li><strong>inference_args</strong> – Non-batchable arguments required as inputs to the model’s
 forward() function. Unlike Tensors in <cite>batch</cite>, these parameters will
 not be dynamically batched</li>
 </ul>
 </td>
 </tr>
 <tr class="field-even field"><th class="field-name">Returns:</th><td class="field-body"><p class="first last">An Iterable of type PredictionResult.</p>
 </td>
 </tr>
 </tbody>
 </table>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.get_num_bytes">
 <code class="descname">get_num_bytes</code><span class="sig-paren">(</span><em>batch: Sequence[&lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9154a0b910&gt;]</em><span class="sig-paren">)</span> &#x2192; int<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.get_num_bytes"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.get_num_bytes" title="Permalink to this definition">¶</a></dt>
 <dd><table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Returns:</th><td class="field-body">The number of bytes of data for a batch of Tensors.</td>
 </tr>
 </tbody>
 </table>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.get_metrics_namespace">
 <code class="descname">get_metrics_namespace</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; str<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.get_metrics_namespace"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.get_metrics_namespace" title="Permalink to this definition">¶</a></dt>
 <dd><table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Returns:</th><td class="field-body">A namespace for metrics collected by the RunInference transform.</td>
 </tr>
 </tbody>
 </table>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.validate_inference_args">
 <code class="descname">validate_inference_args</code><span class="sig-paren">(</span><em>inference_args: Optional[Dict[str, Any]]</em><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.validate_inference_args"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.validate_inference_args" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.batch_elements_kwargs">
 <code class="descname">batch_elements_kwargs</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.batch_elements_kwargs"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.batch_elements_kwargs" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.share_model_across_processes">
 <code class="descname">share_model_across_processes</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; bool<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.share_model_across_processes"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.share_model_across_processes" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.model_copies">
 <code class="descname">model_copies</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; int<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerTensor.model_copies"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerTensor.model_copies" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 </dd></dl>

 <dl class="class">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor">
 <em class="property">class </em><code class="descclassname">apache_beam.ml.inference.pytorch_inference.</code><code class="descname">PytorchModelHandlerKeyedTensor</code><span class="sig-paren">(</span><em>state_dict_path: Optional[str] = None, model_class: Optional[Callable[[...], &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9156cdcc10&gt;]] = None, model_params: Optional[Dict[str, Any]] = None, device: str = 'CPU', *, inference_fn: Callable[[Sequence[Dict[str, &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9157567a90&gt;]], &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9156fe85b0&gt;, &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9156fe8280&gt;, Optional[Dict[str, Any]], Optional[str]], Iterable[apache_beam.ml.inference.base.PredictionResult]] = &lt;function default_keyed_tensor_inference_fn&gt;, torch_script_model_path: Optional[str] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, load_model_args: Optional[Dict[str, Any]] = None, **kwargs</em><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor" title="Permalink to this definition">¶</a></dt>
 <dd><p>Bases: <a class="reference internal" href="apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler" title="apache_beam.ml.inference.base.ModelHandler"><code class="xref py py-class docutils literal notranslate"><span class="pre">apache_beam.ml.inference.base.ModelHandler</span></code></a></p>
 <p>Implementation of the ModelHandler interface for PyTorch.</p>
 <blockquote>
 <div><dl class="docutils">
 <dt>Example Usage for torch model::</dt>
 <dd><dl class="first last docutils">
 <dt>pcoll | RunInference(PytorchModelHandlerKeyedTensor(</dt>
 <dd>state_dict_path=”my_uri”,
 model_class=”my_class”))</dd>
 </dl>
 </dd>
 </dl>
 </div></blockquote>
 <dl class="docutils">
 <dt>Example Usage for torchscript model::</dt>
 <dd><dl class="first last docutils">
 <dt>pcoll | RunInference(PytorchModelHandlerKeyedTensor(</dt>
 <dd>torch_script_model_path=”my_uri”))</dd>
 </dl>
 </dd>
 </dl>
 <p><strong>NOTE:</strong> This API and its implementation are under development and
 do not provide backward compatibility guarantees.</p>
 <p>See <a class="reference external" href="https://pytorch.org/tutorials/beginner/saving_loading_models.html">https://pytorch.org/tutorials/beginner/saving_loading_models.html</a>
 for details</p>
 <table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Parameters:</th><td class="field-body"><ul class="first last simple">
 <li><strong>state_dict_path</strong> – path to the saved dictionary of the model state.</li>
 <li><strong>model_class</strong> – class of the Pytorch model that defines the model
 structure.</li>
 <li><strong>model_params</strong> – A dictionary of arguments required to instantiate the model
 class.</li>
 <li><strong>device</strong> – the device on which you wish to run the model. If
 <code class="docutils literal notranslate"><span class="pre">device</span> <span class="pre">=</span> <span class="pre">GPU</span></code> then a GPU device will be used if it is available.
 Otherwise, it will be CPU.</li>
 <li><strong>inference_fn</strong> – the function to invoke on run_inference.
 default = default_keyed_tensor_inference_fn</li>
 <li><strong>torch_script_model_path</strong> – <dl class="docutils">
 <dt>Path to the torch script model.</dt>
 <dd>the model will be loaded using <cite>torch.jit.load()</cite>.</dd>
 <dt><cite>state_dict_path</cite>, <cite>model_class</cite> and <cite>model_params</cite></dt>
 <dd>arguments will be disregarded.</dd>
 </dl>
 </li>
 <li><strong>min_batch_size</strong> – the minimum batch size to use when batching inputs. This
 batch will be fed into the inference_fn as a Sequence of Keyed Tensors.</li>
 <li><strong>max_batch_size</strong> – the maximum batch size to use when batching inputs. This
 batch will be fed into the inference_fn as a Sequence of Keyed Tensors.</li>
 <li><strong>max_batch_duration_secs</strong> – the maximum amount of time to buffer a batch
 before emitting; used in streaming contexts.</li>
 <li><strong>large_model</strong> – set to true if your model is large enough to run into
 memory pressure if you load multiple copies. Given a model that
 consumes N memory and a machine with W cores and M memory, you should
 set this to True if N*W &gt; M.</li>
 <li><strong>model_copies</strong> – The exact number of models that you would like loaded
 onto your machine. This can be useful if you exactly know your CPU or
 GPU capacity and want to maximize resource utilization.</li>
 <li><strong>load_model_args</strong> – a dictionary of parameters passed to the torch.load
 function to specify custom config for loading models.</li>
 <li><strong>kwargs</strong> – ‘env_vars’ can be used to set environment variables
 before loading the model.</li>
 </ul>
 </td>
 </tr>
 </tbody>
 </table>
 <p><strong>Supported Versions:</strong> RunInference APIs in Apache Beam have been tested
 on torch&gt;=1.9.0,&lt;1.14.0.</p>
 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.load_model">
 <code class="descname">load_model</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9156cdce50&gt;<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.load_model"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.load_model" title="Permalink to this definition">¶</a></dt>
 <dd><p>Loads and initializes a Pytorch model for processing.</p>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.update_model_path">
 <code class="descname">update_model_path</code><span class="sig-paren">(</span><em>model_path: Optional[str] = None</em><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.update_model_path"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.update_model_path" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.run_inference">
 <code class="descname">run_inference</code><span class="sig-paren">(</span><em>batch: Sequence[Dict[str, &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9156cdc8b0&gt;]], model: &lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9156cdc1c0&gt;, inference_args: Optional[Dict[str, Any]] = None</em><span class="sig-paren">)</span> &#x2192; Iterable[apache_beam.ml.inference.base.PredictionResult]<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.run_inference"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.run_inference" title="Permalink to this definition">¶</a></dt>
 <dd><p>Runs inferences on a batch of Keyed Tensors and returns an Iterable of
 Tensor Predictions.</p>
 <p>For the same key across all examples, this will stack all Tensors values
 in a vectorized format to optimize the inference call.</p>
 <table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Parameters:</th><td class="field-body"><ul class="first simple">
 <li><strong>batch</strong> – A sequence of keyed Tensors. These Tensors should be batchable,
 as this method will call <cite>torch.stack()</cite> and pass in batched Tensors
 with dimensions (batch_size, n_features, etc.) into the model’s
 forward() function.</li>
 <li><strong>model</strong> – A PyTorch model.</li>
 <li><strong>inference_args</strong> – Non-batchable arguments required as inputs to the model’s
 forward() function. Unlike Tensors in <cite>batch</cite>, these parameters will
 not be dynamically batched</li>
 </ul>
 </td>
 </tr>
 <tr class="field-even field"><th class="field-name">Returns:</th><td class="field-body"><p class="first last">An Iterable of type PredictionResult.</p>
 </td>
 </tr>
 </tbody>
 </table>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.get_num_bytes">
 <code class="descname">get_num_bytes</code><span class="sig-paren">(</span><em>batch: Sequence[&lt;sphinx.ext.autodoc.importer._MockObject object at 0x7f9156cdc8e0&gt;]</em><span class="sig-paren">)</span> &#x2192; int<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.get_num_bytes"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.get_num_bytes" title="Permalink to this definition">¶</a></dt>
 <dd><table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Returns:</th><td class="field-body">The number of bytes of data for a batch of Dict of Tensors.</td>
 </tr>
 </tbody>
 </table>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.get_metrics_namespace">
 <code class="descname">get_metrics_namespace</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; str<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.get_metrics_namespace"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.get_metrics_namespace" title="Permalink to this definition">¶</a></dt>
 <dd><table class="docutils field-list" frame="void" rules="none">
 <col class="field-name" />
 <col class="field-body" />
 <tbody valign="top">
 <tr class="field-odd field"><th class="field-name">Returns:</th><td class="field-body">A namespace for metrics collected by the RunInference transform.</td>
 </tr>
 </tbody>
 </table>
 </dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.validate_inference_args">
 <code class="descname">validate_inference_args</code><span class="sig-paren">(</span><em>inference_args: Optional[Dict[str, Any]]</em><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.validate_inference_args"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.validate_inference_args" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.batch_elements_kwargs">
 <code class="descname">batch_elements_kwargs</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.batch_elements_kwargs"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.batch_elements_kwargs" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.share_model_across_processes">
 <code class="descname">share_model_across_processes</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; bool<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.share_model_across_processes"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.share_model_across_processes" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 <dl class="method">
 <dt id="apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.model_copies">
 <code class="descname">model_copies</code><span class="sig-paren">(</span><span class="sig-paren">)</span> &#x2192; int<a class="reference internal" href="_modules/apache_beam/ml/inference/pytorch_inference.html#PytorchModelHandlerKeyedTensor.model_copies"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#apache_beam.ml.inference.pytorch_inference.PytorchModelHandlerKeyedTensor.model_copies" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>

 </dd></dl>

 </div>


            </div>

           </div>
           <footer>

     <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">

         <a href="apache_beam.ml.inference.sklearn_inference.html" class="btn btn-neutral float-right" title="apache_beam.ml.inference.sklearn_inference module" accesskey="n" rel="next">Next <span class="fa fa-arrow-circle-right"></span></a>


         <a href="apache_beam.ml.inference.onnx_inference.html" class="btn btn-neutral float-left" title="apache_beam.ml.inference.onnx_inference module" accesskey="p" rel="prev"><span class="fa fa-arrow-circle-left"></span> Previous</a>

     </div>


   <hr/>

   <div role="contentinfo">
     <p>
         &copy; Copyright

     </p>
   </div>
   Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.

 </footer>

         </div>
       </div>

     </section>

   </div>


   <script type="text/javascript">
       jQuery(function () {
           SphinxRtdTheme.Navigation.enable(true);
       });
   </script>


 </body>
 </html>