We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[ONNX] Single input is not exported correctly
When exporting ONNX with single input, input tensor is split and input shape will be wrong.

https://github.com/pfnet/pytorch-pfn-extras/blob/2c5e1440f3137f48d651e705fddb7ea0251583d3/pytorch_pfn_extras/onnx/export_testcase.py#L177

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_pfn_extras/onnx/export_testcase.py`
Content:
```
1 import datetime
2 import io
3 import json
4 import os
5 import subprocess
6 import warnings
7 
8 import onnx
9 import onnx.numpy_helper
10 import torch
11 import torch.autograd
12 from torch.onnx import OperatorExportTypes
13 from torch.onnx.symbolic_helper import _default_onnx_opset_version
14 from torch.onnx.utils import _export as torch_export
15 
16 from pytorch_pfn_extras.onnx.annotate import init_annotate
17 from pytorch_pfn_extras.onnx.strip_large_tensor import \
18     LARGE_TENSOR_DATA_THRESHOLD
19 from pytorch_pfn_extras.onnx.strip_large_tensor import is_large_tensor
20 from pytorch_pfn_extras.onnx.strip_large_tensor import _strip_raw_data
21 from pytorch_pfn_extras.onnx.strip_large_tensor import \
22     _strip_large_initializer_raw_data
23 
24 
25 def _export_meta(model, out_dir, strip_large_tensor_data):
26     ret = {
27         'generated_at': datetime.datetime.now().isoformat(),
28         'output_directory': out_dir,
29         'exporter': 'torch-onnx-utils',
30         'strip_large_tensor_data': strip_large_tensor_data,
31     }
32     try:
33         git_status = subprocess.Popen(['git', 'status'],
34                                       stdout=subprocess.PIPE,
35                                       stderr=subprocess.PIPE)
36         git_status.communicate()
37 
38         def strip_cmd(cmd):
39             return os.popen(cmd).read().strip()
40         if git_status.returncode == 0:
41             ret['git'] = {
42                 'branch': strip_cmd('git rev-parse --abbrev-ref HEAD'),
43                 'commit': strip_cmd('git rev-parse HEAD'),
44                 'remote': strip_cmd('git ls-remote --get-url origin'),
45                 'commit_date': strip_cmd('git show -s --format=%ci HEAD'),
46             }
47     except FileNotFoundError:
48         pass
49 
50     return ret
51 
52 
53 def _export_util(model, args, f, **kwargs):
54     """Wrap operator type to export
55 
56     Copied from torch.onnx.utils.export, to get output values.
57     """
58     aten = kwargs.get('aten', False)
59     export_raw_ir = kwargs.get('export_raw_ir', False)
60     operator_export_type = kwargs.get('operator_export_type', None)
61 
62     if aten or export_raw_ir:
63         assert operator_export_type is None
64         assert aten ^ export_raw_ir
65         operator_export_type = OperatorExportTypes.ATEN if\
66             aten else OperatorExportTypes.RAW
67     elif operator_export_type is None:
68         if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE:
69             operator_export_type = OperatorExportTypes.ONNX_ATEN_FALLBACK
70         else:
71             operator_export_type = OperatorExportTypes.ONNX
72 
73     return torch_export(model, args, f, _retain_param_name=True, **kwargs)
74 
75 
76 def _export(
77         model, args, strip_large_tensor_data=False,
78         large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD, **kwargs):
79     model.zero_grad()
80     bytesio = io.BytesIO()
81     opset_ver = kwargs.get('opset_version', None)
82     if opset_ver is None:
83         opset_ver = _default_onnx_opset_version
84     strip_doc_string = kwargs.pop('strip_doc_string', True)
85     with init_annotate(model, opset_ver) as ann:
86         outs = _export_util(
87             model, args, bytesio, strip_doc_string=False, **kwargs)
88         onnx_graph = onnx.load(io.BytesIO(bytesio.getvalue()))
89         onnx_graph = ann.set_annotate(onnx_graph)
90         onnx_graph = ann.reorg_anchor(onnx_graph)
91     if strip_doc_string:
92         for node in onnx_graph.graph.node:
93             node.doc_string = b''
94     if strip_large_tensor_data:
95         _strip_large_initializer_raw_data(onnx_graph, large_tensor_threshold)
96 
97     return onnx_graph, outs
98 
99 
100 def export(
101         model, args, f, return_output=False, strip_large_tensor_data=False,
102         large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD, **kwargs):
103     """Export model into ONNX Graph.
104 
105     Args:
106         f: A file-like object or a string file path to be written to this
107             file.
108         return_output (bool): If True, return output values come from the
109             model.
110         strip_large_tensor_data (bool): If True, this function will strip
111             data of large tensors to reduce ONNX file size for benchmarking
112         large_tensor_threshold (int): If number of elements of tensor is
113             larger than this value, the tensor is stripped when
114             *strip_large_tensor_data* is True
115     """
116     onnx_graph, outs = _export(
117         model, args, strip_large_tensor_data, large_tensor_threshold,
118         **kwargs)
119 
120     if hasattr(f, 'write'):
121         f.write(onnx_graph.SerializeToString())
122     else:
123         assert isinstance(f, str)
124         warnings.warn(
125             'When export ONNX graph as file, "export_testcase" is '
126             'strongly recommended, please consider use it instead',
127             UserWarning)
128         with open(f, 'wb') as fp:
129             fp.write(onnx_graph.SerializeToString())
130 
131     if return_output:
132         return outs
133 
134 
135 def export_testcase(
136         model, args, out_dir, *, output_grad=False, metadata=True,
137         model_overwrite=True, strip_large_tensor_data=False,
138         large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD,
139         return_output=False, **kwargs):
140     """Export model and I/O tensors of the model in protobuf format.
141 
142     Args:
143         output_grad (bool or Tensor): If True, this function will output
144             model's gradient with names 'gradient_%d.pb'. If set Tensor,
145             use it as gradient *input*. The gradient inputs are output as
146             'gradient_input_%d.pb' along with gradient.
147         metadata (bool): If True, output meta information taken from git log.
148         model_overwrite (bool): If False and model.onnx has already existed,
149             only export input/output data as another test dataset.
150         strip_large_tensor_data (bool): If True, this function will strip
151             data of large tensors to reduce ONNX file size for benchmarking
152         large_tensor_threshold (int): If number of elements of tensor is
153             larger than this value, the tensor is stripped when
154             *strip_large_tensor_data* is True
155     """
156 
157     os.makedirs(out_dir, exist_ok=True)
158     input_names = kwargs.pop(
159         'input_names',
160         ['input_{}'.format(i) for i in range(len(args))])
161     assert len(input_names) == len(args)
162 
163     onnx_graph, outs = _export(
164         model, args, strip_large_tensor_data, large_tensor_threshold,
165         input_names=input_names, **kwargs)
166 
167     # Remove unused inputs
168     # - When keep_initializers_as_inputs=True, inputs contains initializers.
169     #   So we have to filt initializers.
170     # - model.onnx is already issued, so we can modify args here.
171     initializer_names = [init.name for init in onnx_graph.graph.initializer]
172     used_input_index_list = []
173     for used_input in onnx_graph.graph.input:
174         if used_input.name not in initializer_names:
175             used_input_index_list.append(input_names.index(used_input.name))
176     input_names = [input_names[i] for i in used_input_index_list]
177     args = [args[i] for i in used_input_index_list]
178 
179     output_path = os.path.join(out_dir, 'model.onnx')
180     is_on_memory = True
181     if model_overwrite or (not os.path.isfile(output_path)):
182         is_on_memory = False
183         with open(output_path, 'wb') as fp:
184             fp.write(onnx_graph.SerializeToString())
185 
186     def write_to_pb(f, tensor, name=None):
187         array = tensor.detach().cpu().numpy()
188         with open(f, 'wb') as fp:
189             t = onnx.numpy_helper.from_array(array, name)
190             if (strip_large_tensor_data and
191                     is_large_tensor(t, large_tensor_threshold)):
192                 _strip_raw_data(t)
193             fp.write(t.SerializeToString())
194 
195     if isinstance(args, torch.Tensor):
196         args = args,
197     if isinstance(outs, torch.Tensor):
198         outs = outs,
199     data_set_path = os.path.join(out_dir, 'test_data_set_0')
200     seq_id = 0
201     while is_on_memory and os.path.exists(data_set_path):
202         seq_id += 1
203         data_set_path = os.path.join(
204             out_dir, 'test_data_set_{:d}'.format(seq_id))
205     os.makedirs(data_set_path, exist_ok=True)
206     for i, (arg, name) in enumerate(zip(args, input_names)):
207         f = os.path.join(data_set_path, 'input_{}.pb'.format(i))
208         write_to_pb(f, arg, name)
209 
210     output_names = kwargs.get('output_names')
211     if output_names is None:
212         if isinstance(outs, dict):
213             output_names = outs.keys()
214         else:
215             output_names = [None] * len(outs)
216     for i, name in enumerate(output_names):
217         if isinstance(outs, dict):
218             out = outs[name]
219         else:
220             out = outs[i]
221         if isinstance(out, (list, tuple)):
222             assert len(out) == 1, \
223                 'Models returning nested lists/tuples are not supported yet'
224             out = out[0]
225         f = os.path.join(data_set_path, 'output_{}.pb'.format(i))
226         write_to_pb(f, out, name)
227 
228     if output_grad is not False:
229         if isinstance(output_grad, bool):
230             output_grad = [torch.ones_like(outs[idx])
231                            for idx in range(len(output_names))]
232         if isinstance(output_grad, torch.Tensor):
233             output_grad = [output_grad]
234         for idx in range(len(output_names)):
235             write_to_pb(
236                 os.path.join(data_set_path, 'gradient_input_{}.pb'.format(
237                     idx)), output_grad[idx],
238                 output_names[idx])
239         if len(output_names) == len(outs):
240             torch.autograd.backward(outs, grad_tensors=output_grad)
241         else:
242             assert len(
243                 output_names) == 1, 'Single output names is only supported'
244             outs[0].backward(output_grad[0])
245 
246         for i, (name, param) in enumerate(model.named_parameters()):
247             f = os.path.join(data_set_path, 'gradient_{}.pb'.format(i))
248             # NOTE: name does not follow C identifier syntax rules,
249             # like "fc1.bias", not cleanse for now
250             if param.grad is None:
251                 warnings.warn(
252                     'Parameter `{}` does not have gradient value'.format(name))
253             else:
254                 write_to_pb(f, param.grad, name)
255 
256     if metadata:
257         with open(os.path.join(out_dir, 'meta.json'), 'w') as f:
258             json.dump(_export_meta(model, out_dir,
259                                    strip_large_tensor_data), f, indent=2)
260 
261     if return_output:
262         return outs
263
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pytorch_pfn_extras/onnx/export_testcase.py b/pytorch_pfn_extras/onnx/export_testcase.py
--- a/pytorch_pfn_extras/onnx/export_testcase.py
+++ b/pytorch_pfn_extras/onnx/export_testcase.py
@@ -163,6 +163,10 @@
     onnx_graph, outs = _export(
         model, args, strip_large_tensor_data, large_tensor_threshold,
         input_names=input_names, **kwargs)
+    if isinstance(args, torch.Tensor):
+        args = args,
+    if isinstance(outs, torch.Tensor):
+        outs = outs,
 
     # Remove unused inputs
     # - When keep_initializers_as_inputs=True, inputs contains initializers.
@@ -192,10 +196,6 @@
                 _strip_raw_data(t)
             fp.write(t.SerializeToString())
 
-    if isinstance(args, torch.Tensor):
-        args = args,
-    if isinstance(outs, torch.Tensor):
-        outs = outs,
     data_set_path = os.path.join(out_dir, 'test_data_set_0')
     seq_id = 0
     while is_on_memory and os.path.exists(data_set_path):

{"golden_diff": "diff --git a/pytorch_pfn_extras/onnx/export_testcase.py b/pytorch_pfn_extras/onnx/export_testcase.py\n--- a/pytorch_pfn_extras/onnx/export_testcase.py\n+++ b/pytorch_pfn_extras/onnx/export_testcase.py\n@@ -163,6 +163,10 @@\n     onnx_graph, outs = _export(\n         model, args, strip_large_tensor_data, large_tensor_threshold,\n         input_names=input_names, **kwargs)\n+    if isinstance(args, torch.Tensor):\n+        args = args,\n+    if isinstance(outs, torch.Tensor):\n+        outs = outs,\n \n     # Remove unused inputs\n     # - When keep_initializers_as_inputs=True, inputs contains initializers.\n@@ -192,10 +196,6 @@\n                 _strip_raw_data(t)\n             fp.write(t.SerializeToString())\n \n-    if isinstance(args, torch.Tensor):\n-        args = args,\n-    if isinstance(outs, torch.Tensor):\n-        outs = outs,\n     data_set_path = os.path.join(out_dir, 'test_data_set_0')\n     seq_id = 0\n     while is_on_memory and os.path.exists(data_set_path):\n", "issue": "[ONNX] Single input is not exported correctly\nWhen exporting ONNX with single input, input tensor is split and input shape will be wrong.\r\n\r\nhttps://github.com/pfnet/pytorch-pfn-extras/blob/2c5e1440f3137f48d651e705fddb7ea0251583d3/pytorch_pfn_extras/onnx/export_testcase.py#L177\n", "before_files": [{"content": "import datetime\nimport io\nimport json\nimport os\nimport subprocess\nimport warnings\n\nimport onnx\nimport onnx.numpy_helper\nimport torch\nimport torch.autograd\nfrom torch.onnx import OperatorExportTypes\nfrom torch.onnx.symbolic_helper import _default_onnx_opset_version\nfrom torch.onnx.utils import _export as torch_export\n\nfrom pytorch_pfn_extras.onnx.annotate import init_annotate\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import \\\n    LARGE_TENSOR_DATA_THRESHOLD\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import is_large_tensor\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import _strip_raw_data\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import \\\n    _strip_large_initializer_raw_data\n\n\ndef _export_meta(model, out_dir, strip_large_tensor_data):\n    ret = {\n        'generated_at': datetime.datetime.now().isoformat(),\n        'output_directory': out_dir,\n        'exporter': 'torch-onnx-utils',\n        'strip_large_tensor_data': strip_large_tensor_data,\n    }\n    try:\n        git_status = subprocess.Popen(['git', 'status'],\n                                      stdout=subprocess.PIPE,\n                                      stderr=subprocess.PIPE)\n        git_status.communicate()\n\n        def strip_cmd(cmd):\n            return os.popen(cmd).read().strip()\n        if git_status.returncode == 0:\n            ret['git'] = {\n                'branch': strip_cmd('git rev-parse --abbrev-ref HEAD'),\n                'commit': strip_cmd('git rev-parse HEAD'),\n                'remote': strip_cmd('git ls-remote --get-url origin'),\n                'commit_date': strip_cmd('git show -s --format=%ci HEAD'),\n            }\n    except FileNotFoundError:\n        pass\n\n    return ret\n\n\ndef _export_util(model, args, f, **kwargs):\n    \"\"\"Wrap operator type to export\n\n    Copied from torch.onnx.utils.export, to get output values.\n    \"\"\"\n    aten = kwargs.get('aten', False)\n    export_raw_ir = kwargs.get('export_raw_ir', False)\n    operator_export_type = kwargs.get('operator_export_type', None)\n\n    if aten or export_raw_ir:\n        assert operator_export_type is None\n        assert aten ^ export_raw_ir\n        operator_export_type = OperatorExportTypes.ATEN if\\\n            aten else OperatorExportTypes.RAW\n    elif operator_export_type is None:\n        if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE:\n            operator_export_type = OperatorExportTypes.ONNX_ATEN_FALLBACK\n        else:\n            operator_export_type = OperatorExportTypes.ONNX\n\n    return torch_export(model, args, f, _retain_param_name=True, **kwargs)\n\n\ndef _export(\n        model, args, strip_large_tensor_data=False,\n        large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD, **kwargs):\n    model.zero_grad()\n    bytesio = io.BytesIO()\n    opset_ver = kwargs.get('opset_version', None)\n    if opset_ver is None:\n        opset_ver = _default_onnx_opset_version\n    strip_doc_string = kwargs.pop('strip_doc_string', True)\n    with init_annotate(model, opset_ver) as ann:\n        outs = _export_util(\n            model, args, bytesio, strip_doc_string=False, **kwargs)\n        onnx_graph = onnx.load(io.BytesIO(bytesio.getvalue()))\n        onnx_graph = ann.set_annotate(onnx_graph)\n        onnx_graph = ann.reorg_anchor(onnx_graph)\n    if strip_doc_string:\n        for node in onnx_graph.graph.node:\n            node.doc_string = b''\n    if strip_large_tensor_data:\n        _strip_large_initializer_raw_data(onnx_graph, large_tensor_threshold)\n\n    return onnx_graph, outs\n\n\ndef export(\n        model, args, f, return_output=False, strip_large_tensor_data=False,\n        large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD, **kwargs):\n    \"\"\"Export model into ONNX Graph.\n\n    Args:\n        f: A file-like object or a string file path to be written to this\n            file.\n        return_output (bool): If True, return output values come from the\n            model.\n        strip_large_tensor_data (bool): If True, this function will strip\n            data of large tensors to reduce ONNX file size for benchmarking\n        large_tensor_threshold (int): If number of elements of tensor is\n            larger than this value, the tensor is stripped when\n            *strip_large_tensor_data* is True\n    \"\"\"\n    onnx_graph, outs = _export(\n        model, args, strip_large_tensor_data, large_tensor_threshold,\n        **kwargs)\n\n    if hasattr(f, 'write'):\n        f.write(onnx_graph.SerializeToString())\n    else:\n        assert isinstance(f, str)\n        warnings.warn(\n            'When export ONNX graph as file, \"export_testcase\" is '\n            'strongly recommended, please consider use it instead',\n            UserWarning)\n        with open(f, 'wb') as fp:\n            fp.write(onnx_graph.SerializeToString())\n\n    if return_output:\n        return outs\n\n\ndef export_testcase(\n        model, args, out_dir, *, output_grad=False, metadata=True,\n        model_overwrite=True, strip_large_tensor_data=False,\n        large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD,\n        return_output=False, **kwargs):\n    \"\"\"Export model and I/O tensors of the model in protobuf format.\n\n    Args:\n        output_grad (bool or Tensor): If True, this function will output\n            model's gradient with names 'gradient_%d.pb'. If set Tensor,\n            use it as gradient *input*. The gradient inputs are output as\n            'gradient_input_%d.pb' along with gradient.\n        metadata (bool): If True, output meta information taken from git log.\n        model_overwrite (bool): If False and model.onnx has already existed,\n            only export input/output data as another test dataset.\n        strip_large_tensor_data (bool): If True, this function will strip\n            data of large tensors to reduce ONNX file size for benchmarking\n        large_tensor_threshold (int): If number of elements of tensor is\n            larger than this value, the tensor is stripped when\n            *strip_large_tensor_data* is True\n    \"\"\"\n\n    os.makedirs(out_dir, exist_ok=True)\n    input_names = kwargs.pop(\n        'input_names',\n        ['input_{}'.format(i) for i in range(len(args))])\n    assert len(input_names) == len(args)\n\n    onnx_graph, outs = _export(\n        model, args, strip_large_tensor_data, large_tensor_threshold,\n        input_names=input_names, **kwargs)\n\n    # Remove unused inputs\n    # - When keep_initializers_as_inputs=True, inputs contains initializers.\n    #   So we have to filt initializers.\n    # - model.onnx is already issued, so we can modify args here.\n    initializer_names = [init.name for init in onnx_graph.graph.initializer]\n    used_input_index_list = []\n    for used_input in onnx_graph.graph.input:\n        if used_input.name not in initializer_names:\n            used_input_index_list.append(input_names.index(used_input.name))\n    input_names = [input_names[i] for i in used_input_index_list]\n    args = [args[i] for i in used_input_index_list]\n\n    output_path = os.path.join(out_dir, 'model.onnx')\n    is_on_memory = True\n    if model_overwrite or (not os.path.isfile(output_path)):\n        is_on_memory = False\n        with open(output_path, 'wb') as fp:\n            fp.write(onnx_graph.SerializeToString())\n\n    def write_to_pb(f, tensor, name=None):\n        array = tensor.detach().cpu().numpy()\n        with open(f, 'wb') as fp:\n            t = onnx.numpy_helper.from_array(array, name)\n            if (strip_large_tensor_data and\n                    is_large_tensor(t, large_tensor_threshold)):\n                _strip_raw_data(t)\n            fp.write(t.SerializeToString())\n\n    if isinstance(args, torch.Tensor):\n        args = args,\n    if isinstance(outs, torch.Tensor):\n        outs = outs,\n    data_set_path = os.path.join(out_dir, 'test_data_set_0')\n    seq_id = 0\n    while is_on_memory and os.path.exists(data_set_path):\n        seq_id += 1\n        data_set_path = os.path.join(\n            out_dir, 'test_data_set_{:d}'.format(seq_id))\n    os.makedirs(data_set_path, exist_ok=True)\n    for i, (arg, name) in enumerate(zip(args, input_names)):\n        f = os.path.join(data_set_path, 'input_{}.pb'.format(i))\n        write_to_pb(f, arg, name)\n\n    output_names = kwargs.get('output_names')\n    if output_names is None:\n        if isinstance(outs, dict):\n            output_names = outs.keys()\n        else:\n            output_names = [None] * len(outs)\n    for i, name in enumerate(output_names):\n        if isinstance(outs, dict):\n            out = outs[name]\n        else:\n            out = outs[i]\n        if isinstance(out, (list, tuple)):\n            assert len(out) == 1, \\\n                'Models returning nested lists/tuples are not supported yet'\n            out = out[0]\n        f = os.path.join(data_set_path, 'output_{}.pb'.format(i))\n        write_to_pb(f, out, name)\n\n    if output_grad is not False:\n        if isinstance(output_grad, bool):\n            output_grad = [torch.ones_like(outs[idx])\n                           for idx in range(len(output_names))]\n        if isinstance(output_grad, torch.Tensor):\n            output_grad = [output_grad]\n        for idx in range(len(output_names)):\n            write_to_pb(\n                os.path.join(data_set_path, 'gradient_input_{}.pb'.format(\n                    idx)), output_grad[idx],\n                output_names[idx])\n        if len(output_names) == len(outs):\n            torch.autograd.backward(outs, grad_tensors=output_grad)\n        else:\n            assert len(\n                output_names) == 1, 'Single output names is only supported'\n            outs[0].backward(output_grad[0])\n\n        for i, (name, param) in enumerate(model.named_parameters()):\n            f = os.path.join(data_set_path, 'gradient_{}.pb'.format(i))\n            # NOTE: name does not follow C identifier syntax rules,\n            # like \"fc1.bias\", not cleanse for now\n            if param.grad is None:\n                warnings.warn(\n                    'Parameter `{}` does not have gradient value'.format(name))\n            else:\n                write_to_pb(f, param.grad, name)\n\n    if metadata:\n        with open(os.path.join(out_dir, 'meta.json'), 'w') as f:\n            json.dump(_export_meta(model, out_dir,\n                                   strip_large_tensor_data), f, indent=2)\n\n    if return_output:\n        return outs\n", "path": "pytorch_pfn_extras/onnx/export_testcase.py"}], "after_files": [{"content": "import datetime\nimport io\nimport json\nimport os\nimport subprocess\nimport warnings\n\nimport onnx\nimport onnx.numpy_helper\nimport torch\nimport torch.autograd\nfrom torch.onnx import OperatorExportTypes\nfrom torch.onnx.symbolic_helper import _default_onnx_opset_version\nfrom torch.onnx.utils import _export as torch_export\n\nfrom pytorch_pfn_extras.onnx.annotate import init_annotate\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import \\\n    LARGE_TENSOR_DATA_THRESHOLD\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import is_large_tensor\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import _strip_raw_data\nfrom pytorch_pfn_extras.onnx.strip_large_tensor import \\\n    _strip_large_initializer_raw_data\n\n\ndef _export_meta(model, out_dir, strip_large_tensor_data):\n    ret = {\n        'generated_at': datetime.datetime.now().isoformat(),\n        'output_directory': out_dir,\n        'exporter': 'torch-onnx-utils',\n        'strip_large_tensor_data': strip_large_tensor_data,\n    }\n    try:\n        git_status = subprocess.Popen(['git', 'status'],\n                                      stdout=subprocess.PIPE,\n                                      stderr=subprocess.PIPE)\n        git_status.communicate()\n\n        def strip_cmd(cmd):\n            return os.popen(cmd).read().strip()\n        if git_status.returncode == 0:\n            ret['git'] = {\n                'branch': strip_cmd('git rev-parse --abbrev-ref HEAD'),\n                'commit': strip_cmd('git rev-parse HEAD'),\n                'remote': strip_cmd('git ls-remote --get-url origin'),\n                'commit_date': strip_cmd('git show -s --format=%ci HEAD'),\n            }\n    except FileNotFoundError:\n        pass\n\n    return ret\n\n\ndef _export_util(model, args, f, **kwargs):\n    \"\"\"Wrap operator type to export\n\n    Copied from torch.onnx.utils.export, to get output values.\n    \"\"\"\n    aten = kwargs.get('aten', False)\n    export_raw_ir = kwargs.get('export_raw_ir', False)\n    operator_export_type = kwargs.get('operator_export_type', None)\n\n    if aten or export_raw_ir:\n        assert operator_export_type is None\n        assert aten ^ export_raw_ir\n        operator_export_type = OperatorExportTypes.ATEN if\\\n            aten else OperatorExportTypes.RAW\n    elif operator_export_type is None:\n        if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE:\n            operator_export_type = OperatorExportTypes.ONNX_ATEN_FALLBACK\n        else:\n            operator_export_type = OperatorExportTypes.ONNX\n\n    return torch_export(model, args, f, _retain_param_name=True, **kwargs)\n\n\ndef _export(\n        model, args, strip_large_tensor_data=False,\n        large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD, **kwargs):\n    model.zero_grad()\n    bytesio = io.BytesIO()\n    opset_ver = kwargs.get('opset_version', None)\n    if opset_ver is None:\n        opset_ver = _default_onnx_opset_version\n    strip_doc_string = kwargs.pop('strip_doc_string', True)\n    with init_annotate(model, opset_ver) as ann:\n        outs = _export_util(\n            model, args, bytesio, strip_doc_string=False, **kwargs)\n        onnx_graph = onnx.load(io.BytesIO(bytesio.getvalue()))\n        onnx_graph = ann.set_annotate(onnx_graph)\n        onnx_graph = ann.reorg_anchor(onnx_graph)\n    if strip_doc_string:\n        for node in onnx_graph.graph.node:\n            node.doc_string = b''\n    if strip_large_tensor_data:\n        _strip_large_initializer_raw_data(onnx_graph, large_tensor_threshold)\n\n    return onnx_graph, outs\n\n\ndef export(\n        model, args, f, return_output=False, strip_large_tensor_data=False,\n        large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD, **kwargs):\n    \"\"\"Export model into ONNX Graph.\n\n    Args:\n        f: A file-like object or a string file path to be written to this\n            file.\n        return_output (bool): If True, return output values come from the\n            model.\n        strip_large_tensor_data (bool): If True, this function will strip\n            data of large tensors to reduce ONNX file size for benchmarking\n        large_tensor_threshold (int): If number of elements of tensor is\n            larger than this value, the tensor is stripped when\n            *strip_large_tensor_data* is True\n    \"\"\"\n    onnx_graph, outs = _export(\n        model, args, strip_large_tensor_data, large_tensor_threshold,\n        **kwargs)\n\n    if hasattr(f, 'write'):\n        f.write(onnx_graph.SerializeToString())\n    else:\n        assert isinstance(f, str)\n        warnings.warn(\n            'When export ONNX graph as file, \"export_testcase\" is '\n            'strongly recommended, please consider use it instead',\n            UserWarning)\n        with open(f, 'wb') as fp:\n            fp.write(onnx_graph.SerializeToString())\n\n    if return_output:\n        return outs\n\n\ndef export_testcase(\n        model, args, out_dir, *, output_grad=False, metadata=True,\n        model_overwrite=True, strip_large_tensor_data=False,\n        large_tensor_threshold=LARGE_TENSOR_DATA_THRESHOLD,\n        return_output=False, **kwargs):\n    \"\"\"Export model and I/O tensors of the model in protobuf format.\n\n    Args:\n        output_grad (bool or Tensor): If True, this function will output\n            model's gradient with names 'gradient_%d.pb'. If set Tensor,\n            use it as gradient *input*. The gradient inputs are output as\n            'gradient_input_%d.pb' along with gradient.\n        metadata (bool): If True, output meta information taken from git log.\n        model_overwrite (bool): If False and model.onnx has already existed,\n            only export input/output data as another test dataset.\n        strip_large_tensor_data (bool): If True, this function will strip\n            data of large tensors to reduce ONNX file size for benchmarking\n        large_tensor_threshold (int): If number of elements of tensor is\n            larger than this value, the tensor is stripped when\n            *strip_large_tensor_data* is True\n    \"\"\"\n\n    os.makedirs(out_dir, exist_ok=True)\n    input_names = kwargs.pop(\n        'input_names',\n        ['input_{}'.format(i) for i in range(len(args))])\n    assert len(input_names) == len(args)\n\n    onnx_graph, outs = _export(\n        model, args, strip_large_tensor_data, large_tensor_threshold,\n        input_names=input_names, **kwargs)\n    if isinstance(args, torch.Tensor):\n        args = args,\n    if isinstance(outs, torch.Tensor):\n        outs = outs,\n\n    # Remove unused inputs\n    # - When keep_initializers_as_inputs=True, inputs contains initializers.\n    #   So we have to filt initializers.\n    # - model.onnx is already issued, so we can modify args here.\n    initializer_names = [init.name for init in onnx_graph.graph.initializer]\n    used_input_index_list = []\n    for used_input in onnx_graph.graph.input:\n        if used_input.name not in initializer_names:\n            used_input_index_list.append(input_names.index(used_input.name))\n    input_names = [input_names[i] for i in used_input_index_list]\n    args = [args[i] for i in used_input_index_list]\n\n    output_path = os.path.join(out_dir, 'model.onnx')\n    is_on_memory = True\n    if model_overwrite or (not os.path.isfile(output_path)):\n        is_on_memory = False\n        with open(output_path, 'wb') as fp:\n            fp.write(onnx_graph.SerializeToString())\n\n    def write_to_pb(f, tensor, name=None):\n        array = tensor.detach().cpu().numpy()\n        with open(f, 'wb') as fp:\n            t = onnx.numpy_helper.from_array(array, name)\n            if (strip_large_tensor_data and\n                    is_large_tensor(t, large_tensor_threshold)):\n                _strip_raw_data(t)\n            fp.write(t.SerializeToString())\n\n    data_set_path = os.path.join(out_dir, 'test_data_set_0')\n    seq_id = 0\n    while is_on_memory and os.path.exists(data_set_path):\n        seq_id += 1\n        data_set_path = os.path.join(\n            out_dir, 'test_data_set_{:d}'.format(seq_id))\n    os.makedirs(data_set_path, exist_ok=True)\n    for i, (arg, name) in enumerate(zip(args, input_names)):\n        f = os.path.join(data_set_path, 'input_{}.pb'.format(i))\n        write_to_pb(f, arg, name)\n\n    output_names = kwargs.get('output_names')\n    if output_names is None:\n        if isinstance(outs, dict):\n            output_names = outs.keys()\n        else:\n            output_names = [None] * len(outs)\n    for i, name in enumerate(output_names):\n        if isinstance(outs, dict):\n            out = outs[name]\n        else:\n            out = outs[i]\n        if isinstance(out, (list, tuple)):\n            assert len(out) == 1, \\\n                'Models returning nested lists/tuples are not supported yet'\n            out = out[0]\n        f = os.path.join(data_set_path, 'output_{}.pb'.format(i))\n        write_to_pb(f, out, name)\n\n    if output_grad is not False:\n        if isinstance(output_grad, bool):\n            output_grad = [torch.ones_like(outs[idx])\n                           for idx in range(len(output_names))]\n        if isinstance(output_grad, torch.Tensor):\n            output_grad = [output_grad]\n        for idx in range(len(output_names)):\n            write_to_pb(\n                os.path.join(data_set_path, 'gradient_input_{}.pb'.format(\n                    idx)), output_grad[idx],\n                output_names[idx])\n        if len(output_names) == len(outs):\n            torch.autograd.backward(outs, grad_tensors=output_grad)\n        else:\n            assert len(\n                output_names) == 1, 'Single output names is only supported'\n            outs[0].backward(output_grad[0])\n\n        for i, (name, param) in enumerate(model.named_parameters()):\n            f = os.path.join(data_set_path, 'gradient_{}.pb'.format(i))\n            # NOTE: name does not follow C identifier syntax rules,\n            # like \"fc1.bias\", not cleanse for now\n            if param.grad is None:\n                warnings.warn(\n                    'Parameter `{}` does not have gradient value'.format(name))\n            else:\n                write_to_pb(f, param.grad, name)\n\n    if metadata:\n        with open(os.path.join(out_dir, 'meta.json'), 'w') as f:\n            json.dump(_export_meta(model, out_dir,\n                                   strip_large_tensor_data), f, indent=2)\n\n    if return_output:\n        return outs\n", "path": "pytorch_pfn_extras/onnx/export_testcase.py"}]}