|
1 | | -# alexnet |
| 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 | +* features: just several stacked `CRP`(conv-relu-pool) and `CR` layers |
| 9 | +* adaptive average pooling: pytorch can decide its inner parameters, but we need to calculate it ourselves in TensorRT API |
| 10 | +* classifier: just several `fc-relu` layers. All layers can be implemented by tensorrt api, including `addConvolution`, `addActivation`, `addPooling`, `addMatrixMultiply`, `addElementWise` etc. |
8 | 11 |
|
9 | | -``` |
10 | | -// 1. generate alexnet.wts from [pytorchx/alexnet](https://github.com/wang-xinyu/pytorchx/tree/master/alexnet) |
11 | | -
|
12 | | -// 2. put alexnet.wts into tensorrtx/alexnet |
13 | | -
|
14 | | -// 3. build and run |
15 | | -
|
16 | | -cd tensorrtx/alexnet |
| 12 | +## Use AlexNet from PyTorch |
17 | 13 |
|
18 | | -mkdir build |
| 14 | +We can use torchvision to load the pretrained alexnet model: |
19 | 15 |
|
20 | | -cd build |
| 16 | +```python |
| 17 | +alexnet = torchvision.models.alexnet(pretrained=True) |
| 18 | +``` |
21 | 19 |
|
22 | | -cmake .. |
| 20 | +The model structure is: |
| 21 | + |
| 22 | +```txt |
| 23 | +AlexNet( |
| 24 | + (features): Sequential( |
| 25 | + (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2)) |
| 26 | + (1): ReLU(inplace=True) |
| 27 | + (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False) |
| 28 | + (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)) |
| 29 | + (4): ReLU(inplace=True) |
| 30 | + (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False) |
| 31 | + (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) |
| 32 | + (7): ReLU(inplace=True) |
| 33 | + (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) |
| 34 | + (9): ReLU(inplace=True) |
| 35 | + (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) |
| 36 | + (11): ReLU(inplace=True) |
| 37 | + (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False) |
| 38 | + ) |
| 39 | + (avgpool): AdaptiveAvgPool2d(output_size=(6, 6)) |
| 40 | + (classifier): Sequential( |
| 41 | + (0): Dropout(p=0.5, inplace=False) |
| 42 | + (1): Linear(in_features=9216, out_features=4096, bias=True) |
| 43 | + (2): ReLU(inplace=True) |
| 44 | + (3): Dropout(p=0.5, inplace=False) |
| 45 | + (4): Linear(in_features=4096, out_features=4096, bias=True) |
| 46 | + (5): ReLU(inplace=True) |
| 47 | + (6): Linear(in_features=4096, out_features=1000, bias=True) |
| 48 | + ) |
| 49 | +) |
| 50 | +``` |
23 | 51 |
|
24 | | -make |
| 52 | +## FAQ |
25 | 53 |
|
26 | | -sudo ./alexnet -s // serialize model to plan file i.e. 'alexnet.engine' |
| 54 | +### How to align the output with Pytorch? |
27 | 55 |
|
28 | | -sudo ./alexnet -d // deserialize plan file and run inference |
| 56 | +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: |
29 | 57 |
|
30 | | -// 4. see if the output is same as pytorchx/alexnet |
| 58 | +```c++ |
| 59 | +fc3_1->getOutput(0)->setName(OUTPUT_NAME); |
| 60 | +network->markOutput(*pool3->getOutput(0)); // original is: "*fc3_1->getOutput(0)" |
31 | 61 | ``` |
32 | 62 |
|
| 63 | +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. |
| 64 | +
|
| 65 | +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 | 66 |
|
| 67 | +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