class_docstr: End-to-end tests packaging an entire model.
class_docstr_tok: ["End-to-end","tests","packaging","an","entire","model","."]
code: def test_model_save(self):
      # This example shows how you might package a model
      # so that the creator of the model has flexibility about
      # how they want to save it but the 'server' can always
      # use the same API to load the package.
  
      # The convension is for each model to provide a
      # 'model' package with a 'load' function that actual
      # reads the model out of the archive.
  
      # How the load function is implemented is up to the
      # the packager.
  
      # get our normal torchvision resnet
      resnet = resnet18()
  
      f1 = BytesIO()
      # Option 1: save by pickling the whole model
      # + single-line, similar to torch.jit.save
      # - more difficult to edit the code after the model is created
      with PackageExporter(f1) as e:
          e.intern("**")
          e.save_pickle("model", "pickled", resnet)
          # note that this source is the same for all models in this approach
          # so it can be made part of an API that just takes the model and
          # packages it with this source.
          src = dedent(
              """\
                  import importlib
                  import torch_package_importer as resources
  
                  # server knows to call model.load() to get the model,
                  # maybe in the future it passes options as arguments by convension
                  def load():
                      return resources.load_pickle('model', 'pickled')
                  """
          )
          e.save_source_string("model", src, is_package=True)
  
      f2 = BytesIO()
      # Option 2: save with state dict
      # - more code to write to save/load the model
      # + but this code can be edited later to adjust adapt the model later
      with PackageExporter(f2) as e:
          e.intern("**")
          e.save_pickle("model", "state_dict", resnet.state_dict())
          src = dedent(
              """\
                  import importlib
                  import torch_package_importer as resources
  
                  from torchvision.models.resnet import resnet18
                  def load():
                      # if you want, you can later edit how resnet is constructed here
                      # to edit the model in the package, while still loading the original
                      # state dict weights
                      r = resnet18()
                      state_dict = resources.load_pickle('model', 'state_dict')
                      r.load_state_dict(state_dict)
                      return r
                  """
          )
          e.save_source_string("model", src, is_package=True)
  
      # regardless of how we chose to package, we can now use the model in a server in the same way
      input = torch.rand(1, 3, 224, 224)
      results = []
      for m in [f1, f2]:
          m.seek(0)
          importer = PackageImporter(m)
          the_model = importer.import_module("model").load()
          r = the_model(input)
          results.append(r)
  
      self.assertEqual(*results)
  
file_path: pytorch/test/package/test_model.py
imports: from io import BytesIO
  from textwrap import dedent
  from unittest import skipIf
  import torch
  from torch.package import PackageExporter, PackageImporter, sys_importer
  from torch.testing._internal.common_utils import IS_FBCODE, IS_SANDCASTLE, run_tests
label_desc: Use image node_id 2 for calling the ModelTest obj's underlying member method code with example usage: obj.test_model_save() without return types