Torch serve handler doesnt not load state_dict in Docker

venkat.r: <p>I have a saved model which i am trying to serve it with Docker. I can load the stat dicts locally and in Docker if i simply call the <code>trainer_model.pt</code> using <code>load_state_dict()</code>  it works.</p>
<p>However executing it through the <code>handler</code> doesnt seem to we working.</p>
<p>Here is how things are set up.</p>
<h2>
<a name="handerpy-1" class="anchor" href="#handerpy-1"></a><code>hander.py</code>
</h2>
<p>handler which torch-model-archiver needs to read the states.</p>
<pre><code class="lang-auto">import sys
import importlib
from pathlib import Path
sys.path.append("/home/model-server")

import torch
from ts.torch_handler.base_handler import BaseHandler
from ts.utils.util import list_classes_from_module
import torch.nn.functional as FN

print('\n-----Starting handler -----\n')

class DefaultClassifier(BaseHandler):

    def _load_pickeled_model(self, model_dir, model_file, model_pt_path):
        """Load a pickeled model"""

        print(f"Model dir is \n-----&gt; {model_dir} ---------\n")
        print(f"Model file is \n-----&gt; {model_file} ---------\n")


        model_dir = Path(model_dir)
        model_file_path = model_dir / model_file

        if not model_file_path.exists():
            raise RuntimeError(f"Model file {model_file_path} not found")

        module = importlib.import_module(model_file.split(".")[0])
        class_defs_in_module = list_classes_from_module(module)
        if len(class_defs_in_module) != 1:
            raise ValueError(f"Expected one model class, got {class_defs_in_module}")

        model_class = class_defs_in_module[0]
        model = model_class()

        if model_pt_path:
            print(f"Model pt path is \n-----&gt; {model_pt_path} ---------\n")
            model.load_state_dict(torch.load(model_pt_path)['model'], strict=False)

        return model
</code></pre>
<h2>
<a name="docker-file-2" class="anchor" href="#docker-file-2"></a>Docker file</h2>
<p>Container in which the model is serving.</p>
<pre><code class="lang-auto">FROM pytorch/torchserve:0.4.2-cpu

USER root
RUN printf "\nservice_envelope=json" &gt;&gt; /home/model-server/config.properties \
    &amp;&amp; pip install pytorch-lightning torchmetrics google-cloud-storage scikit-learn
USER model-server


COPY src /home/model-server/src
COPY load_model.py /home/model-server/src
COPY configs /home/model-server/configs
COPY /models/trained-model-2021-11-23--23:08-hs.pt /home/model-server/src/

RUN torch-model-archiver \
    --model-name=ponzi \
    --version=0.1 \
    --model-file=/home/model-server/src/model.py \
    --serialized-file=/home/model-server/src/trained-model-2021-11-23--23:08-hs.pt \
    --handler=/home/model-server/src/handler.py \
    --export-path=/home/model-server/model-store \
    --extra-files=/home/model-server/configs/features.yml,/home/model-server/configs/model_configs.yml,/home/model-server/src/utils.py

CMD ["torchserve", \
     "--start", \
     "--ts-config=/home/model-server/config.properties", \
     "--models ponzi=ponzi.mar"]
</code></pre>
<p>For some reason hander starts but never really calls <code>_load_pickeled_model</code> as i see it from the logs.</p>
<pre><code class="lang-auto">2021-11-24 10:28:14,156 [INFO ] W-9008-ponzi_v20211122-stdout MODEL_LOG - Connection accepted: /home/model-server/tmp/.ts.sock.9008.
2021-11-24 10:28:14,169 [INFO ] W-9006-ponzi_v20211122-stdout MODEL_LOG - model_name: ponzi, batchSize: 1
2021-11-24 10:28:14,170 [INFO ] W-9008-ponzi_v20211122-stdout MODEL_LOG - model_name: ponzi, batchSize: 1
2021-11-24 10:28:14,171 [INFO ] W-9006-ponzi_v20211122-stdout MODEL_LOG - 
2021-11-24 10:28:14,171 [INFO ] W-9006-ponzi_v20211122-stdout MODEL_LOG - -----Starting handler -----
2021-11-24 10:28:14,172 [INFO ] W-9008-ponzi_v20211122-stdout MODEL_LOG - 
2021-11-24 10:28:14,172 [INFO ] W-9008-ponzi_v20211122-stdout MODEL_LOG - -----Starting handler -----
2021-11-24 10:28:14,203 [INFO ] W-9009-ponzi_v20211122-stdout MODEL_LOG - model_name: ponzi, batchSize: 1
2021-11-24 10:28:14,204 [INFO ] W-9009-ponzi_v20211122-stdout MODEL_LOG - 
2021-11-24 10:28:14,204 [INFO ] W-9009-ponzi_v20211122-stdout MODEL_LOG - -----Starting handler -----
2021-11-24 10:28:15,323 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG - 
2021-11-24 10:28:15,323 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG - Backend worker process died.
2021-11-24 10:28:15,323 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG - Traceback (most recent call last):
2021-11-24 10:28:15,323 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/ts/model_service_worker.py", line 183, in &lt;module&gt;
2021-11-24 10:28:15,323 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     worker.run_server()
2021-11-24 10:28:15,323 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/ts/model_service_worker.py", line 155, in run_server
2021-11-24 10:28:15,324 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     self.handle_connection(cl_socket)
2021-11-24 10:28:15,324 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/ts/model_service_worker.py", line 117, in handle_connection
2021-11-24 10:28:15,324 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     service, result, code = self.load_model(msg)
2021-11-24 10:28:15,324 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/ts/model_service_worker.py", line 90, in load_model
2021-11-24 10:28:15,324 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     service = model_loader.load(model_name, model_dir, handler, gpu, batch_size, envelope)
2021-11-24 10:28:15,325 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/ts/model_loader.py", line 110, in load
2021-11-24 10:28:15,325 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     initialize_fn(service.context)
2021-11-24 10:28:15,325 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/ts/torch_handler/base_handler.py", line 66, in initialize
2021-11-24 10:28:15,325 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     self.model = self._load_pickled_model(model_dir, model_file, model_pt_path)
2021-11-24 10:28:15,325 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/ts/torch_handler/base_handler.py", line 130, in _load_pickled_model
2021-11-24 10:28:15,325 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     model.load_state_dict(state_dict)
2021-11-24 10:28:15,326 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -   File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 1407, in load_state_dict
2021-11-24 10:28:15,326 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG -     self.__class__.__name__, "\n\t".join(error_msgs)))
2021-11-24 10:28:15,326 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG - RuntimeError: Error(s) in loading state_dict for LitPonzi:
2021-11-24 10:28:15,326 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG - 	Missing key(s) in state_dict: "layer_1.weight", "layer_1.bias", "layer_3.weight", "layer_3.bias". 
2021-11-24 10:28:15,326 [INFO ] W-9010-ponzi_v20211122-stdout MODEL_LOG - 	Unexpected key(s) in state_dict: "model". 
2021-11-24 10:28:15,327 [INFO ] epollEventLoopGroup-5-13 org.pytorch.serve.wlm.WorkerThread - 9010 Worker disconnected. WORKER_STARTED
2021-11-24 10:28:15,328 [DEBUG] W-9010-ponzi_v20211122 org.pytorch.serve.wlm.WorkerThread - System state is : WORKER_STARTED
2021-11-24 10:28:15,328 [DEBUG] W-9010-ponzi_v20211122 org.pytorch.serve.wlm.WorkerThread - Backend worker monitoring thread interrupted or backend worker process died.
</code></pre>
<p>I have tested loading the <code>state_dict</code> as standalone in the container and that works. For example the below works</p>
<pre><code class="lang-auto">import torch
from src.model import LitPonzi

# test to be run in docker
loaded_model = LitPonzi()
loaded_model.load_state_dict(torch.load('src/trained-model-2021-11-23--23:08-hs.pt')['model'], strict=False)
loaded_model.eval()
print(f"Loaded model, {loaded_model}")
</code></pre>
<p>Any idea what i am doing wrong here.?  Appreciate your help.</p>

marksaroufim: <p>I think the fix is just <code>jit.load()</code> here’s an examplee you can follow that does this</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/serve/blob/master/examples/Huggingface_Transformers/Transformer_handler_generalized.py#L63">
  <header class="source">

      <a href="https://github.com/pytorch/serve/blob/master/examples/Huggingface_Transformers/Transformer_handler_generalized.py#L63" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/serve/blob/master/examples/Huggingface_Transformers/Transformer_handler_generalized.py#L63" target="_blank" rel="noopener nofollow ugc">pytorch/serve/blob/master/examples/Huggingface_Transformers/Transformer_handler_generalized.py#L63</a></h4>



    <pre class="onebox">      <code class="lang-py">
        <ol class="start lines" start="53" style="counter-reset: li-counter 52 ;">
            <li>    logger.warning("Missing the setup_config.json file.")</li>
            <li>
            </li>
<li># Loading the shared object of compiled Faster Transformer Library if Faster Transformer is set</li>
            <li>if self.setup_config["FasterTransformer"]:</li>
            <li>    faster_transformer_complied_path = os.path.join(model_dir, "libpyt_fastertransformer.so")</li>
            <li>    torch.classes.load_library(faster_transformer_complied_path)</li>
            <li>
            </li>
<li># Loading the model and tokenizer from checkpoint and config files based on the user's choice of mode</li>
            <li># further setup config can be added.</li>
            <li>if self.setup_config["save_mode"] == "torchscript":</li>
            <li class="selected">    self.model = torch.jit.load(model_pt_path, map_location=self.device)</li>
            <li>elif self.setup_config["save_mode"] == "pretrained":</li>
            <li>    if self.setup_config["mode"] == "sequence_classification":</li>
            <li>        self.model = AutoModelForSequenceClassification.from_pretrained(</li>
            <li>            model_dir</li>
            <li>        )</li>
            <li>    elif self.setup_config["mode"] == "question_answering":</li>
            <li>        self.model = AutoModelForQuestionAnswering.from_pretrained(model_dir)</li>
            <li>    elif self.setup_config["mode"] == "token_classification":</li>
            <li>        self.model = AutoModelForTokenClassification.from_pretrained(model_dir)</li>
            <li>    else:</li>
        </ol>
      </code>
    </pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


engineero: <p><a class="mention" href="/u/venkat.r">@venkat.r</a>, did <a class="mention" href="/u/marksaroufim">@marksaroufim</a>’s fix work for you? I am experiencing a similar problem and not finding a lot of resources yet.</p>