Skip to content

Commit b28afb9

Browse files
authored
Update README.md
1 parent f174029 commit b28afb9

File tree

1 file changed

+8
-8
lines changed

1 file changed

+8
-8
lines changed

README.md

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,19 +12,19 @@ If you find an issue, please [let us know](../..//issues)!
1212
> Please note, this converter has limited coverage of TensorRT / PyTorch. We created it primarily
1313
> to easily optimize the models used in the [JetBot](https://github.com/NVIDIA-AI-IOT/jetbot) project. If you find the converter helpful with other models, please [let us know](../..//issues).
1414
15-
### Setup
15+
## Setup
1616

1717
```bash
1818
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
1919
cd torch2trt
2020
python setup.py install --user
2121
```
2222

23-
### Usage
23+
## Usage
2424

2525
Below are some usage examples, for more check out the [notebooks](notebooks).
2626

27-
#### Convert
27+
### Convert
2828

2929
```python
3030
from torch2trt import torch2trt
@@ -40,7 +40,7 @@ x = torch.ones((1, 3, 224, 224)).cuda()
4040
model_trt = torch2trt(model, [x])
4141
```
4242

43-
#### Execute
43+
### Execute
4444

4545
We can execute returned ``TRTModule`` just like the original PyTorch model
4646

@@ -52,7 +52,7 @@ y_trt = model_trt(x)
5252
print(torch.max(torch.abs(y - y_trt)))
5353
```
5454

55-
#### Save and load
55+
### Save and load
5656

5757
We can save the model as a ``state_dict``.
5858

@@ -70,7 +70,7 @@ model_trt = TRTModule()
7070
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))
7171
```
7272

73-
### Models
73+
## Models
7474

7575
We tested the converter against these models using the [test.sh](test.sh) script. You can generate the results by calling
7676

@@ -104,7 +104,7 @@ bash test.sh TEST_OUTPUT.md
104104
| vgg19_bn | | | 51.4 | 121 |
105105

106106

107-
### How does it work?
107+
## How does it work?
108108

109109
This converter works by attaching conversion functions (like ``convert_ReLU``) to the original
110110
PyTorch functional calls (like ``torch.nn.ReLU.forward``). The sample input data is passed
@@ -117,7 +117,7 @@ uses this ``_trt`` to add layers to the TensorRT network, and then sets the ``_t
117117
relevant output tensors. Once the model is fully executed, the final tensors returns are marked as outputs
118118
of the TensorRT network, and the optimized TensorRT engine is built.
119119

120-
### How to add (or override) a converter
120+
## How to add (or override) a converter
121121

122122
Here we show how to add a converter for the ``ReLU`` module using the TensorRT
123123
python API.

0 commit comments

Comments
 (0)