|
1 | 1 | # alexnet |
2 | 2 |
|
3 | | -AlexNet model architecture from the "One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper. |
| 3 | +## Introduction |
4 | 4 |
|
5 | | -For the details, you can refer to [pytorchx/alexnet](https://github.com/wang-xinyu/pytorchx/tree/master/alexnet) |
| 5 | +AlexNet model architecture comes from this paper: [One weird trick for parallelizing convolutional neural networks](https://arxiv.org/abs/1404.5997). To generate `.wts` file, you can refer to [pytorchx/alexnet](https://github.com/wang-xinyu/pytorchx/tree/master/alexnet). To check the pytorch implementation of AlexNet, refer to [HERE](https://github.com/pytorch/vision/blob/main/torchvision/models/alexnet.py#L17) |
6 | 6 |
|
7 | | -This alexnet is just several `conv-relu-pool` blocks followed by several `fc-relu`, nothing special. All layers can be implemented by tensorrt api, including `addConvolution`, `addActivation`, `addPooling`, `addFullyConnected`. |
| 7 | +AlexNet consists of 3 major parts: features, adaptive average pooling, and classifier: |
8 | 8 |
|
| 9 | +- features: just several stacked `CRP`(conv-relu-pool) and `CR` layers |
| 10 | +- adaptive average pooling: pytorch can decide its inner parameters, but we need to calculate it ourselves in TensorRT API |
| 11 | +- classifier: just several `fc-relu` layers. All layers can be implemented by tensorrt api, including `addConvolution`, `addActivation`, `addPooling`, `addMatrixMultiply`, `addElementWise` etc. |
| 12 | + |
| 13 | +## Use AlexNet from PyTorch |
| 14 | + |
| 15 | +We can use torchvision to load the pretrained alexnet model: |
| 16 | + |
| 17 | +```python |
| 18 | +alexnet = torchvision.models.alexnet(pretrained=True) |
9 | 19 | ``` |
10 | | -// 1. generate alexnet.wts from [pytorchx/alexnet](https://github.com/wang-xinyu/pytorchx/tree/master/alexnet) |
11 | 20 |
|
12 | | -// 2. put alexnet.wts into tensorrtx/alexnet |
| 21 | +The model structure is: |
| 22 | + |
| 23 | +```bash |
| 24 | +AlexNet( |
| 25 | + (features): Sequential( |
| 26 | + (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2)) |
| 27 | + (1): ReLU(inplace=True) |
| 28 | + (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False) |
| 29 | + (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)) |
| 30 | + (4): ReLU(inplace=True) |
| 31 | + (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False) |
| 32 | + (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) |
| 33 | + (7): ReLU(inplace=True) |
| 34 | + (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) |
| 35 | + (9): ReLU(inplace=True) |
| 36 | + (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) |
| 37 | + (11): ReLU(inplace=True) |
| 38 | + (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False) |
| 39 | + ) |
| 40 | + (avgpool): AdaptiveAvgPool2d(output_size=(6, 6)) |
| 41 | + (classifier): Sequential( |
| 42 | + (0): Dropout(p=0.5, inplace=False) |
| 43 | + (1): Linear(in_features=9216, out_features=4096, bias=True) |
| 44 | + (2): ReLU(inplace=True) |
| 45 | + (3): Dropout(p=0.5, inplace=False) |
| 46 | + (4): Linear(in_features=4096, out_features=4096, bias=True) |
| 47 | + (5): ReLU(inplace=True) |
| 48 | + (6): Linear(in_features=4096, out_features=1000, bias=True) |
| 49 | + ) |
| 50 | +) |
| 51 | +``` |
13 | 52 |
|
14 | | -// 3. build and run |
| 53 | +## Usage |
15 | 54 |
|
16 | | -cd tensorrtx/alexnet |
| 55 | +1. use `gen_wts.py` to generate wts file. |
17 | 56 |
|
18 | | -mkdir build |
| 57 | +```bash |
| 58 | +python3 gen_wts.py |
| 59 | +``` |
19 | 60 |
|
20 | | -cd build |
| 61 | +2. build C++ code |
21 | 62 |
|
22 | | -cmake .. |
| 63 | +```bash |
| 64 | +pushd tensorrtx/alexnet |
| 65 | +cmake -S . -B build -G Ninja --fresh |
| 66 | +cmake --build build |
| 67 | +``` |
23 | 68 |
|
24 | | -make |
| 69 | +3. serialize wts model to engine file. |
25 | 70 |
|
26 | | -sudo ./alexnet -s // serialize model to plan file i.e. 'alexnet.engine' |
| 71 | +```bash |
| 72 | +./build/alexnet -s |
| 73 | +``` |
27 | 74 |
|
28 | | -sudo ./alexnet -d // deserialize plan file and run inference |
| 75 | +4. run inference |
29 | 76 |
|
30 | | -// 4. see if the output is same as pytorchx/alexnet |
| 77 | +```bash |
| 78 | +./build/alexnet -d |
31 | 79 | ``` |
32 | 80 |
|
| 81 | +output looks like: |
| 82 | + |
| 83 | +```txt |
| 84 | +... |
| 85 | +==== |
| 86 | +Execution time: 1ms |
| 87 | +0.1234, -0.5678, ... |
| 88 | +==== |
| 89 | +prediction result: |
| 90 | +Top: 0 idx: 285, logits: 9.9, label: Egyptian cat |
| 91 | +Top: 1 idx: 281, logits: 8.304, label: tabby, tabby cat |
| 92 | +Top: 2 idx: 282, logits: 6.859, label: tiger cat |
| 93 | +``` |
| 94 | + |
| 95 | +## FAQ |
| 96 | + |
| 97 | +### How to align the output with Pytorch? |
| 98 | + |
| 99 | +If your output is different from pytorch, you have to check which TensorRT API or your code cause this. A simple solution would be check the `.engine` output part by part, e.g., you can set the early layer of alexnet as output: |
| 100 | + |
| 101 | +```c++ |
| 102 | +fc3_1->getOutput(0)->setName(OUTPUT_NAME); |
| 103 | +network->markOutput(*pool3->getOutput(0)); // original is: "*fc3_1->getOutput(0)" |
| 104 | +``` |
| 105 | +
|
| 106 | +For this line of code, i use the output from "feature" part of alexnet, ignoring the rest of the model, then, don't forget to change the `OUTPUT_SIZE` macro on top of the file, lastly, build the `.engine` file to apply the changes. |
| 107 | +
|
| 108 | +You can sum up all output from C++ code, and compare it with Pytorch output, for Pytorch, you can do this by: `torch.sum(x)` at debug phase. The ideal value deviation between 2 values would be $[10^{-1}, 10^{-2}]$, for this example, since the output elements for "feature" is $256 * 6 * 6$ (bacth = 1), the final error would roughly be $10^{-4}$. |
33 | 109 |
|
| 110 | +Note: This is a quick check, for more accurate check, you have to save the output tensor into a file to compare them value by value, but this situation is rare. |
0 commit comments