├── .gitattributes
├── .gitignore
├── LICENSE
├── README.md
├── calibrator.py
├── engine_function.py
├── model_function.py
├── train.py
├── val_tensorrt.py
└── val_tensorrt_sim.py


/.gitattributes:
--------------------------------------------------------------------------------
1 | # Auto detect text files and perform LF normalization
2 | * text=auto
3 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Byte-compiled / optimized / DLL files
  2 | __pycache__/
  3 | *.py[cod]
  4 | *$py.class
  5 | 
  6 | # C extensions
  7 | *.so
  8 | 
  9 | # Distribution / packaging
 10 | .Python
 11 | build/
 12 | develop-eggs/
 13 | dist/
 14 | downloads/
 15 | eggs/
 16 | .eggs/
 17 | lib/
 18 | lib64/
 19 | parts/
 20 | sdist/
 21 | var/
 22 | wheels/
 23 | *.egg-info/
 24 | .installed.cfg
 25 | *.egg
 26 | MANIFEST
 27 | 
 28 | # PyInstaller
 29 | #  Usually these files are written by a python script from a template
 30 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 31 | *.manifest
 32 | *.spec
 33 | 
 34 | # Installer logs
 35 | pip-log.txt
 36 | pip-delete-this-directory.txt
 37 | 
 38 | # Unit test / coverage reports
 39 | htmlcov/
 40 | .tox/
 41 | .nox/
 42 | .coverage
 43 | .coverage.*
 44 | .cache
 45 | nosetests.xml
 46 | coverage.xml
 47 | *.cover
 48 | .hypothesis/
 49 | .pytest_cache/
 50 | 
 51 | # Translations
 52 | *.mo
 53 | *.pot
 54 | 
 55 | # Django stuff:
 56 | *.log
 57 | local_settings.py
 58 | db.sqlite3
 59 | 
 60 | # Flask stuff:
 61 | instance/
 62 | .webassets-cache
 63 | 
 64 | # Scrapy stuff:
 65 | .scrapy
 66 | 
 67 | # Sphinx documentation
 68 | docs/_build/
 69 | 
 70 | # PyBuilder
 71 | target/
 72 | 
 73 | # Jupyter Notebook
 74 | .ipynb_checkpoints
 75 | 
 76 | # IPython
 77 | profile_default/
 78 | ipython_config.py
 79 | 
 80 | # pyenv
 81 | .python-version
 82 | 
 83 | # celery beat schedule file
 84 | celerybeat-schedule
 85 | 
 86 | # SageMath parsed files
 87 | *.sage.py
 88 | 
 89 | # Environments
 90 | .env
 91 | .venv
 92 | env/
 93 | venv/
 94 | ENV/
 95 | env.bak/
 96 | venv.bak/
 97 | 
 98 | # Spyder project settings
 99 | .spyderproject
100 | .spyproject
101 | 
102 | # Rope project settings
103 | .ropeproject
104 | 
105 | # mkdocs documentation
106 | /site
107 | 
108 | # mypy
109 | .mypy_cache/
110 | .dmypy.json
111 | dmypy.json
112 | 
113 | # Pyre type checker
114 | .pyre/
115 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2019 darkknightzh
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # TensorRT_pytorch
 2 | 
 3 | A simple demo to train mnist in pytorch and speed up inference by TensorRT.
 4 | The training code comes from [here](https://github.com/pytorch/examples/tree/master/mnist) .
 5 | The code to use TensorRT comes from samples in installation package of TensorRT.
 6 | 
 7 | The TensorRT version we use is 5.1.5. The code may not compatible with other versions of TensorRT.
 8 | 
 9 | You can go to [（原）pytorch中使用TensorRT](https://www.cnblogs.com/darkknightzh/p/11332155.html) for more information.
10 | 


--------------------------------------------------------------------------------
/calibrator.py:
--------------------------------------------------------------------------------
  1 | #
  2 | # Copyright 1993-2019 NVIDIA Corporation.  All rights reserved.
  3 | #
  4 | # NOTICE TO LICENSEE:
  5 | #
  6 | # This source code and/or documentation ("Licensed Deliverables") are
  7 | # subject to NVIDIA intellectual property rights under U.S. and
  8 | # international Copyright laws.
  9 | #
 10 | # These Licensed Deliverables contained herein is PROPRIETARY and
 11 | # CONFIDENTIAL to NVIDIA and is being provided under the terms and
 12 | # conditions of a form of NVIDIA software license agreement by and
 13 | # between NVIDIA and Licensee ("License Agreement") or electronically
 14 | # accepted by Licensee.  Notwithstanding any terms or conditions to
 15 | # the contrary in the License Agreement, reproduction or disclosure
 16 | # of the Licensed Deliverables to any third party without the express
 17 | # written consent of NVIDIA is prohibited.
 18 | #
 19 | # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
 20 | # LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
 21 | # SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE.  IT IS
 22 | # PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
 23 | # NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
 24 | # DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
 25 | # NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
 26 | # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
 27 | # LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
 28 | # SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
 29 | # DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
 30 | # WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
 31 | # ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
 32 | # OF THESE LICENSED DELIVERABLES.
 33 | #
 34 | # U.S. Government End Users.  These Licensed Deliverables are a
 35 | # "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
 36 | # 1995), consisting of "commercial computer software" and "commercial
 37 | # computer software documentation" as such terms are used in 48
 38 | # C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
 39 | # only as a commercial end item.  Consistent with 48 C.F.R.12.212 and
 40 | # 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
 41 | # U.S. Government End Users acquire the Licensed Deliverables with
 42 | # only those rights set forth herein.
 43 | #
 44 | # Any use of the Licensed Deliverables in individual and commercial
 45 | # software must include, in the user documentation and internal
 46 | # comments to the code, the above Disclaimer and U.S. Government End
 47 | # Users Notice.
 48 | 
 49 | import tensorrt as trt
 50 | import os
 51 | import json
 52 | 
 53 | import pycuda.driver as cuda
 54 | import pycuda.autoinit
 55 | from PIL import Image
 56 | import numpy as np
 57 | 
 58 | # For reading size information from batches
 59 | import struct
 60 | import torch
 61 | import cv2
 62 | import math
 63 | 
 64 | class MNISTEntropyCalibrator(trt.IInt8EntropyCalibrator2):
 65 |     def __init__(self, batch_data_dir, cache_file):
 66 |         # Whenever you specify a custom constructor for a TensorRT class,
 67 |         # you MUST call the constructor of the parent explicitly.
 68 |         trt.IInt8EntropyCalibrator2.__init__(self)
 69 | 
 70 |         self.cache_file = cache_file
 71 |         # Get a list of all the batch files in the batch folder.
 72 |         self.batch_files = [os.path.join(batch_data_dir, f) for f in os.listdir(batch_data_dir)]
 73 | 
 74 |         # Find out the shape of a batch and then allocate a device buffer of that size.
 75 |         self.shape, _, _ = self.read_batch_file(self.batch_files[0])
 76 |         # Each element of the calibration data is a float32.
 77 |         self.device_input = cuda.mem_alloc(trt.volume(self.shape) * trt.float32.itemsize)
 78 | 
 79 |         # Create a generator that will give us batches. We can use next() to iterate over the result.
 80 |         def load_batches():
 81 |             for f in self.batch_files:
 82 |                 shape, data, labels = self.read_batch_file(f)
 83 |                 yield shape, data, labels
 84 |         self.batches = load_batches()
 85 | 
 86 |     # This function is used to load calibration data from the calibration batch files.
 87 |     # In this implementation, one file corresponds to one batch, but it is also possible to use
 88 |     # aggregate data from multiple files, or use only data from portions of a file.
 89 |     def read_batch_file(self, filename):
 90 |         with open(filename, "rb") as f:
 91 |             # Read the first 4 integers. These will be the NCHW dimensions of the data.
 92 |             shape = tuple(struct.unpack("<L", f.read(trt.int32.itemsize))[0] for _ in range(4))
 93 |             # Next read in all the images, where each element of each image is a float32.
 94 |             data = f.read(trt.volume(shape) * trt.float32.itemsize)
 95 |             # The remainder of the file consists of labels
 96 |             labels = f.read()
 97 |         return shape, data, labels
 98 | 
 99 |     def get_batch_size(self):
100 |         return self.shape[0]
101 | 
102 |     # TensorRT passes along the names of the engine bindings to the get_batch function.
103 |     # You don't necessarily have to use them, but they can be useful to understand the order of
104 |     # the inputs. The bindings list is expected to have the same ordering as 'names'.
105 |     def get_batch(self, names):
106 |         try:
107 |             # Get a single batch.
108 |             _, data, _ = next(self.batches)
109 |             # Copy to device, then return a list containing pointers to input device buffers.
110 |             cuda.memcpy_htod(self.device_input, data)
111 |             return [int(self.device_input)]
112 |         except StopIteration:
113 |             # When we're out of batches, we return either [] or None.
114 |             # This signals to TensorRT that there is no calibration data remaining.
115 |             return None
116 | 
117 |     def read_calibration_cache(self):
118 |         # If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.
119 |         if os.path.exists(self.cache_file):
120 |             with open(self.cache_file, "rb") as f:
121 |                 return f.read()
122 | 
123 |     def write_calibration_cache(self, cache):
124 |         with open(self.cache_file, "wb") as f:
125 |             f.write(cache)
126 | 
127 | 
128 | class ExampleEntropyCalibrator(trt.IInt8EntropyCalibrator2):
129 |     def __init__(self, datafolder, cache_file, c, h, w):
130 |         # Whenever you specify a custom constructor for a TensorRT class,
131 |         # you MUST call the constructor of the parent explicitly.
132 |         trt.IInt8EntropyCalibrator2.__init__(self)
133 | 
134 |         self.cache_file = cache_file
135 | 
136 |         data, targets = torch.load(datafolder)
137 |         self.all_files = data.numpy()
138 | 
139 |         # Find out the shape of a batch and then allocate a device buffer of that size.
140 |         # self.shape, _, _ = self.read_batch_file(self.batch_files[0])
141 |         self.shape = [1, c, h, w]
142 |         # Each element of the calibration data is a float32.
143 |         self.device_input = cuda.mem_alloc(trt.volume(self.shape) * trt.float32.itemsize)
144 | 
145 |         # Create a generator that will give us batches. We can use next() to iterate over the result.
146 |         def load_batches():
147 |             for idx in range(len(self.all_files)):
148 |                 data = self.read_batch_file(idx)
149 |                 yield data
150 |         self.batches = load_batches()
151 | 
152 |     # This function is used to load calibration data from the calibration batch files.
153 |     # In this implementation, one file corresponds to one batch, but it is also possible to use
154 |     # aggregate data from multiple files, or use only data from portions of a file.
155 |     def read_batch_file(self, idx):
156 |         def normalize(input, mean, std):
157 |             input = input / 255.0
158 |             return (input - mean) / std
159 | 
160 |         img = self.all_files[idx]
161 |         normed_img = normalize(img, mean=(0.1307,), std=(0.3081,))
162 |         normed_img = normed_img.astype(np.float32)
163 |         return normed_img
164 | 
165 |     def get_batch_size(self):
166 |         return self.shape[0]
167 | 
168 |     # TensorRT passes along the names of the engine bindings to the get_batch function.
169 |     # You don't necessarily have to use them, but they can be useful to understand the order of
170 |     # the inputs. The bindings list is expected to have the same ordering as 'names'.
171 |     def get_batch(self, names):
172 |         try:
173 |             # Get a single batch.
174 |             data = next(self.batches)
175 |             # Copy to device, then return a list containing pointers to input device buffers.
176 |             cuda.memcpy_htod(self.device_input, data)
177 |             return [int(self.device_input)]
178 |         except StopIteration:
179 |             # When we're out of batches, we return either [] or None.
180 |             # This signals to TensorRT that there is no calibration data remaining.
181 |             print('exception')
182 |             return None
183 | 
184 |     def read_calibration_cache(self):
185 |         # If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.
186 |         if os.path.exists(self.cache_file):
187 |             with open(self.cache_file, "rb") as f:
188 |                 return f.read()
189 | 
190 |     def write_calibration_cache(self, cache):
191 |         with open(self.cache_file, "wb") as f:
192 |             f.write(cache)


--------------------------------------------------------------------------------
/engine_function.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import tensorrt as trt
 3 | import pycuda.driver as cuda
 4 | import pycuda.autoinit    # 此句代码中未使用，但是必须有。this is useful, otherwise stream = cuda.Stream() will cause 'explicit_context_dependent failed: invalid device context - no currently active context?'
 5 | 
 6 | class HostDeviceMem(object):
 7 |     def __init__(self, host_mem, device_mem):
 8 |         self.host = host_mem
 9 |         self.device = device_mem
10 | 
11 |     def __str__(self):
12 |         return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
13 | 
14 |     def __repr__(self):
15 |         return self.__str__()
16 | 
17 | def allocate_buffers(engine):
18 |     inputs = []
19 |     outputs = []
20 |     bindings = []
21 |     stream = cuda.Stream()
22 |     for binding in engine:
23 |         dtype = trt.nptype(engine.get_binding_dtype(binding))
24 |         # Allocate host and device buffers
25 |         host_mem = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size, dtype)
26 |         device_mem = cuda.mem_alloc(host_mem.nbytes)
27 |         # Append the device buffer to device bindings.
28 |         bindings.append(int(device_mem))
29 |         # Append to the appropriate list.
30 |         if engine.binding_is_input(binding):
31 |             inputs.append(HostDeviceMem(host_mem, device_mem))
32 |         else:
33 |             outputs.append(HostDeviceMem(host_mem, device_mem))
34 |     return inputs, outputs, bindings, stream
35 | 
36 | 
37 | def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
38 |     # Transfer input data to the GPU.
39 |     [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
40 |     # Run inference.
41 |     context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
42 |     # Transfer predictions back from the GPU.
43 |     [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
44 |     # Synchronize the stream
45 |     stream.synchronize()
46 |     # Return only the host outputs.
47 |     return [out.host for out in outputs]
48 | 
49 | 
50 | def saveONNX(model, filepath, c, h, w):
51 |     model = model.cuda()
52 |     dummy_input = torch.randn(1, c, h, w, device='cuda')
53 |     torch.onnx.export(model, dummy_input, filepath, verbose=True)
54 | 
55 | 
56 | def build_engine(onnx_file_path):
57 |     TRT_LOGGER = trt.Logger(trt.Logger.WARNING)   # INFO
58 |     # For more information on TRT basics, refer to the introductory samples.
59 |     with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
60 |         if builder.platform_has_fast_fp16:
61 |             print('this card support fp16')
62 |         if builder.platform_has_fast_int8:
63 |             print('this card support int8')
64 | 
65 |         builder.max_workspace_size = 1 << 30
66 |         with open(onnx_file_path, 'rb') as model:
67 |            parser.parse(model.read())
68 |         return builder.build_cuda_engine(network)
69 | 
70 | # This function builds an engine from a Caffe model.
71 | def build_engine_int8(onnx_file_path, calib):
72 |     TRT_LOGGER = trt.Logger()
73 |     with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
74 |         # We set the builder batch size to be the same as the calibrator's, as we use the same batches
75 |         # during inference. Note that this is not required in general, and inference batch size is
76 |         # independent of calibration batch size.
77 |         builder.max_batch_size = 1  # calib.get_batch_size()
78 |         builder.max_workspace_size = 1 << 30
79 |         builder.int8_mode = True
80 |         builder.int8_calibrator = calib
81 |         with open(onnx_file_path, 'rb') as model:
82 |            parser.parse(model.read())   # , dtype=trt.float32
83 |         return builder.build_cuda_engine(network)
84 | 
85 | def save_engine(engine, engine_dest_path):
86 |     buf = engine.serialize()
87 |     with open(engine_dest_path, 'wb') as f:
88 |         f.write(buf)
89 | 
90 | def load_engine(engine_path):
91 |     TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  # INFO
92 |     with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
93 |         return runtime.deserialize_cuda_engine(f.read())
94 | 


--------------------------------------------------------------------------------
/model_function.py:
--------------------------------------------------------------------------------
 1 | from torch import nn
 2 | 
 3 | class Example_Model(nn.Module):
 4 |     def __init__(self):
 5 |         super().__init__()
 6 |         self.model1 = nn.Conv2d(3, 24, 3)
 7 |         self.ReLU = nn.ReLU()
 8 |         self.model2 = nn.Conv2d(24, 3, 3)
 9 | 
10 |     def forward(self, x):
11 |         out = self.ReLU(self.model1(x))
12 |         out = self.model2(out)
13 |         return out
14 | 


--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
  1 | from __future__ import print_function
  2 | import argparse
  3 | import torch
  4 | import torch.nn as nn
  5 | import torch.nn.functional as F
  6 | import torch.optim as optim
  7 | from torchvision import datasets, transforms
  8 | 
  9 | 
 10 | class Net(nn.Module):
 11 |     def __init__(self):
 12 |         super(Net, self).__init__()
 13 |         self.conv1 = nn.Conv2d(1, 20, 5, 1)
 14 |         self.conv2 = nn.Conv2d(20, 50, 5, 1)
 15 |         self.fc1 = nn.Linear(4 * 4 * 50, 500)
 16 |         self.fc2 = nn.Linear(500, 10)
 17 | 
 18 |     def forward(self, x):
 19 |         x = F.relu(self.conv1(x))
 20 |         x = F.max_pool2d(x, 2, 2)
 21 |         x = F.relu(self.conv2(x))
 22 |         x = F.max_pool2d(x, 2, 2)
 23 |         x = x.view(-1, 4 * 4 * 50)
 24 |         x = F.relu(self.fc1(x))
 25 |         x = self.fc2(x)
 26 |         return F.log_softmax(x, dim=1)
 27 | 
 28 | 
 29 | def train(args, model, device, train_loader, optimizer, epoch):
 30 |     model.train()
 31 |     for batch_idx, (data, target) in enumerate(train_loader):
 32 |         data, target = data.to(device), target.to(device)
 33 |         optimizer.zero_grad()
 34 |         output = model(data)
 35 |         loss = F.nll_loss(output, target)
 36 |         loss.backward()
 37 |         optimizer.step()
 38 |         if batch_idx % args.log_interval == 0:
 39 |             print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
 40 |                 epoch, batch_idx * len(data), len(train_loader.dataset),
 41 |                        100. * batch_idx / len(train_loader), loss.item()))
 42 | 
 43 | 
 44 | def test(args, model, device, test_loader):
 45 |     model.eval()
 46 |     test_loss = 0
 47 |     correct = 0
 48 |     with torch.no_grad():
 49 |         for data, target in test_loader:
 50 |             data, target = data.to(device), target.to(device)
 51 |             output = model(data)
 52 |             test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
 53 |             pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
 54 |             correct += pred.eq(target.view_as(pred)).sum().item()
 55 | 
 56 |     test_loss /= len(test_loader.dataset)
 57 | 
 58 |     print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
 59 |         test_loss, correct, len(test_loader.dataset),
 60 |         100. * correct / len(test_loader.dataset)))
 61 | 
 62 | 
 63 | def main():
 64 |     # Training settings
 65 |     parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
 66 |     parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)')
 67 |     parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', help='input batch size for testing (default: 1000)')
 68 |     parser.add_argument('--epochs', type=int, default=3, metavar='N', help='number of epochs to train (default: 10)')
 69 |     parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)')
 70 |     parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)')
 71 |     parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training')
 72 |     parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)')
 73 |     parser.add_argument('--log-interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status')
 74 | 
 75 |     parser.add_argument('--save-model', action='store_true', default=True, help='For Saving the current Model')
 76 |     args = parser.parse_args()
 77 |     use_cuda = not args.no_cuda and torch.cuda.is_available()
 78 | 
 79 |     torch.manual_seed(args.seed)
 80 | 
 81 |     device = torch.device("cuda" if use_cuda else "cpu")
 82 | 
 83 |     kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
 84 |     train_loader = torch.utils.data.DataLoader(
 85 |         datasets.MNIST('data', train=True, download=True,
 86 |                        transform=transforms.Compose([
 87 |                            transforms.ToTensor(),
 88 |                            transforms.Normalize((0.1307,), (0.3081,))
 89 |                        ])),
 90 |         batch_size=args.batch_size, shuffle=True, **kwargs)
 91 |     test_loader = torch.utils.data.DataLoader(
 92 |         datasets.MNIST('data', train=False, transform=transforms.Compose([
 93 |             transforms.ToTensor(),
 94 |             transforms.Normalize((0.1307,), (0.3081,))
 95 |         ])),
 96 |         batch_size=args.test_batch_size, shuffle=False, **kwargs)
 97 | 
 98 |     model = Net().to(device)
 99 |     optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
100 | 
101 |     for epoch in range(1, args.epochs + 1):
102 |         train(args, model, device, train_loader, optimizer, epoch)
103 |         test(args, model, device, test_loader)
104 |         if (args.save_model):
105 |             torch.save(model.state_dict(), "mnist_cnn_{}.pt".format(epoch))
106 | 
107 | 
108 | if __name__ == '__main__':
109 |     main()


--------------------------------------------------------------------------------
/val_tensorrt.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import numpy as np
  3 | import sys
  4 | from PIL import Image
  5 | 
  6 | #from model_function import Example_Model
  7 | from train import Net
  8 | from engine_function import *
  9 | import tensorrt as trt
 10 | import pycuda.driver as cuda
 11 | import pycuda.autoinit    # 此句代码中未使用，但是必须有。this is useful, otherwise stream = cuda.Stream() will cause 'explicit_context_dependent failed: invalid device context - no currently active context?'
 12 | 
 13 | import calibrator
 14 | import collections
 15 | 
 16 | def normalize(input, mean, std):
 17 |     input = input/255.0
 18 |     return (input - mean)/std
 19 | 
 20 | 
 21 | def test_pytorch():
 22 |     net = Net()
 23 |     checkpoint_path = 'mnist_cnn_3.pt'
 24 |     net.load_state_dict(torch.load(checkpoint_path))
 25 | 
 26 |     correct = 0
 27 |     data, targets = torch.load('data/MNIST/processed/test.pt')
 28 |     data, targets = data.numpy(), targets.numpy()
 29 |     with torch.no_grad():
 30 |         for index in range(len(data)):
 31 |             if index%100 == 0:
 32 |                 print(index, len(data))
 33 |             img, target = data[index], int(targets[index])
 34 |             img = normalize(img, mean=(0.1307,), std=(0.3081,))
 35 |             img = img[np.newaxis, np.newaxis, :,:]
 36 |             img = torch.from_numpy(img)
 37 |             output = net(img.float())
 38 |             pred = np.argmax(output.numpy(), axis=1)
 39 |             correct += pred[0] == target
 40 | 
 41 |     print('\nAccuracy: {}  ({:.0f}%)\n'.format(correct, 100. * correct / len(targets)))
 42 | 
 43 | 
 44 | def export_onnx():
 45 |     net = Net()
 46 |     checkpoint_path = 'mnist_cnn_3.pt'
 47 |     net.load_state_dict(torch.load(checkpoint_path))
 48 |     saveONNX(net, 'mnist_cnn_3.onnx', 1, 28, 28)
 49 | 
 50 | 
 51 | def test_onnx_fp32():
 52 |     #torch.load('check_point.pth')   # 如果结果不对，加上这句话
 53 |     with build_engine('mnist_cnn_3.onnx') as engine, engine.create_execution_context() as context:
 54 |         #save_engine(engine, "mnist_cnn_3_fp32.engine")
 55 |         inputs, outputs, bindings, stream = allocate_buffers(engine)
 56 | 
 57 |         correct = 0
 58 |         data, targets = torch.load('data/MNIST/processed/test.pt')
 59 |         data, targets = data.numpy(), targets.numpy()
 60 |         for index in range(len(data)):
 61 |             if index % 100 == 0:
 62 |                 print(index, len(data))
 63 |             img, target = data[index], int(targets[index])
 64 |             img = normalize(img, mean=(0.1307,), std=(0.3081,))
 65 | 
 66 |             img_numpy = img.ravel().astype(np.float32)
 67 |             np.copyto(inputs[0].host, img_numpy)
 68 |             output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
 69 |             output = [np.reshape(stage_i, (10)) for stage_i in output]  # 有多个输出时遍历
 70 | 
 71 |             pred = np.argmax(output, axis=1)
 72 |             correct += pred[0] == target
 73 | 
 74 |         print('\nAccuracy: {}  ({:.0f}%)\n'.format(correct, 100. * correct / len(targets)))
 75 | 
 76 | def test_onnx_fp32_engine():
 77 |     #torch.load('check_point.pth')   # 如果结果不对，加上这句话
 78 |     with load_engine("mnist_cnn_3_fp32.engine") as engine, engine.create_execution_context() as context:
 79 |         inputs, outputs, bindings, stream = allocate_buffers(engine)
 80 | 
 81 |         correct = 0
 82 |         data, targets = torch.load('data/MNIST/processed/test.pt')
 83 |         data, targets = data.numpy(), targets.numpy()
 84 |         for index in range(len(data)):
 85 |             if index % 100 == 0:
 86 |                 print(index, len(data))
 87 |             img, target = data[index], int(targets[index])
 88 |             img = normalize(img, mean=(0.1307,), std=(0.3081,))
 89 | 
 90 |             img_numpy = img.ravel().astype(np.float32)
 91 |             np.copyto(inputs[0].host, img_numpy)
 92 |             output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
 93 |             output = [np.reshape(stage_i, (10)) for stage_i in output]  # 有多个输出时遍历
 94 | 
 95 |             pred = np.argmax(output, axis=1)
 96 |             correct += pred[0] == target
 97 | 
 98 |         print('\nAccuracy: {}  ({:.0f}%)\n'.format(correct, 100. * correct / len(targets)))
 99 | 
100 | 
101 | def test_onnx_int8():
102 |     #torch.load('mnist_cnn_3.pth')    # 如果结果不对，加上这句话
103 |     calibration_cache = "mnist_cnn_3.cache"
104 |     calib = calibrator.ExampleEntropyCalibrator(datafolder='data/MNIST/processed/test.pt', cache_file=calibration_cache, c=1, h=28, w=28)
105 | 
106 |     with build_engine_int8('mnist_cnn_3.onnx', calib) as engine, engine.create_execution_context() as context:
107 |         #save_engine(engine, "mnist_cnn_3_int8.engine")
108 |         inputs, outputs, bindings, stream = allocate_buffers(engine)  # 最好在下面do_inference之前分配一次内存即可
109 | 
110 |         correct = 0
111 |         data, targets = torch.load('data/MNIST/processed/test.pt')
112 |         data, targets = data.numpy(), targets.numpy()
113 |         for index in range(len(data)):
114 |             if index % 100 == 0:
115 |                 print(index, len(data))
116 |             img, target = data[index], int(targets[index])
117 |             img = normalize(img, mean=(0.1307,), std=(0.3081,))
118 | 
119 |             img_numpy = img.ravel().astype(np.float32)
120 |             np.copyto(inputs[0].host, img_numpy)
121 |             output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
122 |             output = [np.reshape(stage_i, (10)) for stage_i in output]  # 有多个输出时遍历
123 | 
124 |             pred = np.argmax(output, axis=1)
125 |             correct += pred[0] == target
126 | 
127 |         print('\nAccuracy: {}  ({:.0f}%)\n'.format(correct, 100. * correct / len(targets)))
128 |         
129 | 
130 | def test_onnx_int8_engine():
131 |     #torch.load('mnist_cnn_3.pth')
132 |     with load_engine("mnist_cnn_3_int8.engine") as engine, engine.create_execution_context() as context:
133 |         inputs, outputs, bindings, stream = allocate_buffers(engine)   # 最好在下面do_inference之前分配一次内存即可
134 | 
135 |         correct = 0
136 |         data, targets = torch.load('data/MNIST/processed/test.pt')
137 |         data, targets = data.numpy(), targets.numpy()
138 |         for index in range(len(data)):
139 |             if index % 100 == 0:
140 |                 print(index, len(data))
141 |             img, target = data[index], int(targets[index])
142 |             img = normalize(img, mean=(0.1307,), std=(0.3081,))
143 | 
144 |             img_numpy = img.ravel().astype(np.float32)
145 |             np.copyto(inputs[0].host, img_numpy)
146 |             output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
147 |             output = [np.reshape(stage_i, (10)) for stage_i in output]  # 有多个输出时遍历
148 | 
149 |             pred = np.argmax(output, axis=1)
150 |             correct += pred[0] == target
151 | 
152 |         print('\nAccuracy: {}  ({:.0f}%)\n'.format(correct, 100. * correct / len(targets)))
153 | 
154 | 
155 | def main():
156 |     test_pytorch()
157 |     #export_onnx()
158 |     #test_onnx_fp32()
159 |     #test_onnx_fp32_engine()
160 |     #test_onnx_int8()
161 |     #test_onnx_int8_engine()
162 | 
163 | 
164 | if __name__ == '__main__':
165 |     main()
166 | 
167 | 


--------------------------------------------------------------------------------
/val_tensorrt_sim.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import numpy as np
  3 | import sys
  4 | 
  5 | from model_function import Example_Model
  6 | from engine_function import *
  7 | import tensorrt as trt
  8 | import pycuda.driver as cuda
  9 | import pycuda.autoinit    # 此句代码中未使用，但是必须有。this is useful, otherwise stream = cuda.Stream() will cause 'explicit_context_dependent failed: invalid device context - no currently active context?'
 10 | 
 11 | import calibrator
 12 | import collections
 13 | 
 14 | 
 15 | def parse_arguments(argv):
 16 |     parser = argparse.ArgumentParser()
 17 |     args = parser.parse_args()
 18 |     return args
 19 | 
 20 | 
 21 | def load_state(net, checkpoint):
 22 |     source_state = checkpoint['state_dict']
 23 |     target_state = net.state_dict()
 24 |     new_target_state = collections.OrderedDict()
 25 |     for target_key, target_value in target_state.items():
 26 |         if target_key in source_state and source_state[target_key].size() == target_state[target_key].size():
 27 |             new_target_state[target_key] = source_state[target_key]
 28 |         else:
 29 |             new_target_state[target_key] = target_state[target_key]
 30 |             print('[WARNING] Not found pre-trained parameters for {}'.format(target_key))
 31 | 
 32 |     net.load_state_dict(new_target_state)
 33 | 
 34 | def test_pytorch(args):
 35 |     net = Example_Model()
 36 | 
 37 |     checkpoint_path = 'check_point.pth'
 38 |     #torch.save({'state_dict': net.state_dict()}, checkpoint_path)
 39 | 
 40 |     checkpoint = torch.load(checkpoint_path)
 41 |     load_state(net, checkpoint)
 42 | 
 43 |     #input = torch.randn(1, 3, 24, 24)
 44 |     #np.save('input.npy', input.cpu())
 45 |     input = np.load('input.npy')
 46 |     input = torch.tensor(input)
 47 | 
 48 |     print(input)
 49 | 
 50 |     out = net(input)
 51 |     print(out)
 52 | 
 53 | 
 54 | 
 55 | def export_onnx(args):
 56 |     net = Example_Model()
 57 | 
 58 |     checkpoint_path = 'check_point.pth'
 59 |     checkpoint = torch.load(checkpoint_path)
 60 |     load_state(net, checkpoint)
 61 |     saveONNX(net, 'check_point.onnx')
 62 | 
 63 | 
 64 | def test_onnx_fp32(args):
 65 |     #torch.load('check_point.pth')   # 如果结果不对，加上这句话
 66 |     with build_engine('check_point.onnx') as engine, engine.create_execution_context() as context:
 67 |         #save_engine(engine, "check_point_fp32.engine")
 68 |         inputs, outputs, bindings, stream = allocate_buffers(engine)
 69 |         img_numpy = np.load('input.npy')
 70 |         img_numpy = img_numpy.ravel().astype(np.float32)
 71 |         np.copyto(inputs[0].host, img_numpy)
 72 |         output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
 73 |         output = [np.reshape(stage_i, (-1, 20, 20)) for stage_i in output]  # 有多个输出时遍历
 74 |         print(output)
 75 | 
 76 | def test_onnx_fp32_engine(args):
 77 |     #torch.load('check_point.pth')   # 如果结果不对，加上这句话
 78 |     with load_engine("check_point_fp32.engine") as engine, engine.create_execution_context() as context:
 79 |         inputs, outputs, bindings, stream = allocate_buffers(engine)
 80 |         img_numpy = np.load('input.npy')
 81 |         img_numpy = img_numpy.ravel().astype(np.float32)
 82 |         np.copyto(inputs[0].host, img_numpy)
 83 |         output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
 84 |         output = [np.reshape(stage_i, (-1, 20, 20)) for stage_i in output]  # 有多个输出时遍历
 85 |         print(output)
 86 | 
 87 | 
 88 | def test_onnx_int8(args):   # 此处未测试，因为没有数据，正常测试可以
 89 |     torch.load('check_point.pth')    # 如果结果不对，加上这句话
 90 |     calibration_cache = "check_point.cache"
 91 |     calib = calibrator.ExampleEntropyCalibrator(datafolder='your_data_folder',
 92 |         file_list='your_json_file.json', cache_file=calibration_cache)
 93 | 
 94 |     with build_engine_int8('check_point.onnx', calib) as engine, engine.create_execution_context() as context:
 95 |         save_engine(engine, "check_point_int8.engine")
 96 |         inputs, outputs, bindings, stream = allocate_buffers(engine)  # 最好在下面do_inference之前分配一次内存即可
 97 |         img_numpy = np.load('input.npy')
 98 |         img_numpy = img_numpy.ravel().astype(np.float32)
 99 |         np.copyto(inputs[0].host, img_numpy)
100 |         output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
101 |         output = [np.reshape(stage_i, (-1, 20, 20)) for stage_i in output]  # 有多个输出时遍历
102 |         print(output)
103 | 
104 | def test_onnx_int8_engine(args):
105 |     torch.load('check_point.pth')
106 |     with load_engine("check_point_int8.engine") as engine, engine.create_execution_context() as context:
107 |         inputs, outputs, bindings, stream = allocate_buffers(engine)   # 最好在下面do_inference之前分配一次内存即可
108 |         img_numpy = np.load('input.npy')
109 |         img_numpy = img_numpy.ravel().astype(np.float32)
110 |         np.copyto(inputs[0].host, img_numpy)
111 |         output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
112 |         output = [np.reshape(stage_i, (-1, 20, 20)) for stage_i in output]  # 有多个输出时遍历
113 |         print(output)
114 | 
115 | 
116 | def main(args):
117 |     test_pytorch(args)
118 |     #export_onnx(args)
119 |     #test_onnx_fp32(args)
120 |     #test_onnx_fp32_engine(args)
121 |     #test_onnx_int8(args)
122 |     #test_onnx_int8_engine(args)
123 | 
124 | 
125 | if __name__ == '__main__':
126 |     main(parse_arguments(sys.argv[1:]))
127 | 
128 | 


--------------------------------------------------------------------------------