mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
5bcbbf5373
Summary: *Context:* https://github.com/pytorch/pytorch/issues/53406 added a lint for trailing whitespace at the ends of lines. However, in order to pass FB-internal lints, that PR also had to normalize the trailing newlines in four of the files it touched. This PR adds an OSS lint to normalize trailing newlines. The changes to the following files (made in 54847d0adb9be71be4979cead3d9d4c02160e4cd) are the only manually-written parts of this PR: - `.github/workflows/lint.yml` - `mypy-strict.ini` - `tools/README.md` - `tools/test/test_trailing_newlines.py` - `tools/trailing_newlines.py` I would have liked to make this just a shell one-liner like the other three similar lints, but nothing I could find quite fit the bill. Specifically, all the answers I tried from the following Stack Overflow questions were far too slow (at least a minute and a half to run on this entire repository): - [How to detect file ends in newline?](https://stackoverflow.com/q/38746) - [How do I find files that do not end with a newline/linefeed?](https://stackoverflow.com/q/4631068) - [How to list all files in the Git index without newline at end of file](https://stackoverflow.com/q/27624800) - [Linux - check if there is an empty line at the end of a file [duplicate]](https://stackoverflow.com/q/34943632) - [git ensure newline at end of each file](https://stackoverflow.com/q/57770972) To avoid giving false positives during the few days after this PR is merged, we should probably only merge it after https://github.com/pytorch/pytorch/issues/54967. Pull Request resolved: https://github.com/pytorch/pytorch/pull/54737 Test Plan: Running the shell script from the "Ensure correct trailing newlines" step in the `quick-checks` job of `.github/workflows/lint.yml` should print no output and exit in a fraction of a second with a status of 0. That was not the case prior to this PR, as shown by this failing GHA workflow run on an earlier draft of this PR: - https://github.com/pytorch/pytorch/runs/2197446987?check_suite_focus=true In contrast, this run (after correcting the trailing newlines in this PR) succeeded: - https://github.com/pytorch/pytorch/pull/54737/checks?check_run_id=2197553241 To unit-test `tools/trailing_newlines.py` itself (this is run as part of our "Test tools" GitHub Actions workflow): ``` python tools/test/test_trailing_newlines.py ``` Reviewed By: malfet Differential Revision: D27409736 Pulled By: samestep fbshipit-source-id: 46f565227046b39f68349bbd5633105b2d2e9b19
109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
## @package onnx
|
|
#Module caffe2.python.trt.transform
|
|
|
|
"""
|
|
TensorRT related transformation
|
|
Note that ONNX-TRT enforce an NCHW input!
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from caffe2.proto import caffe2_pb2
|
|
from caffe2.python import workspace
|
|
import caffe2.python._import_c_extension as C
|
|
import numpy as np
|
|
|
|
def _dim_values_to_list(dim_values):
|
|
return [x.dim_value for x in dim_values]
|
|
|
|
|
|
def _get_output_shapes(output_value_infos):
|
|
names = [x.name for x in output_value_infos]
|
|
shapes = [_dim_values_to_list(x.type.tensor_type.shape.dim) for x in output_value_infos]
|
|
return dict(zip(names, shapes))
|
|
|
|
|
|
def check_gpu_():
|
|
try:
|
|
C.get_cuda_version()
|
|
except Exception as _:
|
|
raise Exception("TensorRT related functions require CUDA support")
|
|
|
|
def convert_onnx_model_to_trt_op(onnx_model,
|
|
max_batch_size=64,
|
|
max_workspace_size=2*1024*1024,
|
|
verbosity=1,
|
|
debug_builder=False):
|
|
"""
|
|
Convert the whole ONNX model to a TensorRT C2 op
|
|
"""
|
|
check_gpu_()
|
|
trt_str = C.onnx_to_trt_op(onnx_model.SerializeToString(),
|
|
_get_output_shapes(onnx_model.graph.output),
|
|
max_batch_size,
|
|
max_workspace_size,
|
|
verbosity,
|
|
debug_builder)
|
|
op = caffe2_pb2.OperatorDef()
|
|
op.ParseFromString(trt_str)
|
|
return op
|
|
|
|
|
|
# Assume the workspace is already filled with init weights
|
|
def _infer_shapes(pred_net, inputs):
|
|
workspace.RunNetOnce(pred_net)
|
|
hints = {}
|
|
for op in pred_net.op:
|
|
for o in op.output:
|
|
if o not in hints:
|
|
blob = workspace.FetchBlob(o)
|
|
if hasattr(blob, 'shape'):
|
|
hints[o] = blob.shape
|
|
for i in op.input:
|
|
if i not in hints:
|
|
blob = workspace.FetchBlob(i)
|
|
if hasattr(blob, 'shape'):
|
|
hints[i] = blob.shape
|
|
|
|
return hints
|
|
|
|
|
|
def transform_caffe2_net(
|
|
pred_net,
|
|
input_shapes,
|
|
populate_shapes = False,
|
|
max_batch_size=64,
|
|
max_workspace_size=2*1024*1024,
|
|
verbosity=1,
|
|
debug_builder=False,
|
|
build_serializable_op=True):
|
|
"""
|
|
Transform the caffe2_net by collapsing TRT-runnable nodes into trt c2 ops
|
|
"""
|
|
check_gpu_()
|
|
|
|
# Hacky way to infer shapes as not all our operators have shape inference function.
|
|
# Normally this is not needed
|
|
shape_hints = {}
|
|
if populate_shapes:
|
|
input_data = {}
|
|
for k,v in input_shapes.items():
|
|
input_data[k] = np.random.randn(*v).astype(np.float32)
|
|
shape_hints = _infer_shapes(pred_net, input_data)
|
|
|
|
for k,v in input_shapes.items():
|
|
shape_hints[k] = v
|
|
pred_net_str = C.transform_trt(pred_net.SerializeToString(),
|
|
shape_hints,
|
|
max_batch_size,
|
|
max_workspace_size,
|
|
verbosity,
|
|
debug_builder,
|
|
build_serializable_op)
|
|
pred_net_cut = caffe2_pb2.NetDef()
|
|
pred_net_cut.ParseFromString(pred_net_str)
|
|
return pred_net_cut
|