code: def main(args):
      os.environ["PADDLESEG_EXPORT_STAGE"] = "True"
      cfg = Config(args.config)
      builder = SegBuilder(cfg)
      net = builder.model
  
      if args.model_path is not None:
          para_state_dict = paddle.load(args.model_path)
          net.set_dict(para_state_dict)
          logger.info("Loaded trained params of model successfully.")
  
      if args.input_shape is None:
          shape = [None, 3, None, None]
      else:
          shape = args.input_shape
  
      output_op = args.output_op
      if args.without_argmax:
          logger.warning(
              "--without_argmax will be deprecated, please use --output_op"
          )
          output_op = "none"
      if args.with_softmax:
          logger.warning(
              "--with_softmax will be deprecated, please use --output_op"
          )
          output_op = "softmax"
  
      new_net = (
          net
          if output_op == "none"
          else SavedSegmentationNet(net, output_op)
      )
      new_net.eval()
      new_net = paddle.jit.to_static(
          new_net,
          input_spec=[
              paddle.static.InputSpec(shape=shape, dtype="float32")
          ],
      )
  
      save_path = os.path.join(args.save_dir, "model")
      paddle.jit.save(new_net, save_path)
  
      yml_file = os.path.join(args.save_dir, "deploy.yaml")
      with open(yml_file, "w") as file:
          transforms = cfg.val_dataset_cfg.get(
              "transforms", [{"type": "Normalize"}]
          )
          output_dtype = "int32" if output_op == "argmax" else "float32"
          data = {
              "Deploy": {
                  "model": "model.pdmodel",
                  "params": "model.pdiparams",
                  "transforms": transforms,
                  "input_shape": shape,
                  "output_op": output_op,
                  "output_dtype": output_dtype,
              }
          }
          yaml.dump(data, file)
  
      logger.info(f"The inference model is saved in {args.save_dir}")
  
file_path: PaddleSeg/contrib/CityscapesSOTA/export.py
imports: import argparse
  import os
  import paddle
  import yaml
  from paddleseg.cvlibs import Config, SegBuilder
  from paddleseg.utils import logger
  import datasets, models
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types