本专栏主要是深度学习/自动驾驶相关的源码实现,获取全套代码请参考
实现一个回归模型
数据生成
假设输入为shape=(10,5)的input_tensor，真值为shape=(1,1)的gt。
def gen_data(length):
data_dicts = []
for _ in range(length):
input = np.random.rand(10, 5)
w = np.array([[1], [0.5], [1], [-1], [1]])
b = np.array([[0]])
output = input.dot(w) + b
output = np.max(output, axis=0)
data_dict = {}
data_dict['input'] = input
data_dict['output'] = output
data_dicts.append(data_dict)
return data_dicts
该函数gen_data生成一个长度为length的数据集，其中每个数据项是一个字典，包含随机生成的10x5输入矩阵input和一个基于该输入和固定权重w、偏置b计算得到的输出值output（取为线性变换后的最大值）。返回的数据集是一个包含这些字典的列表。
模型搭建
class MModel(torch.nn.Module):
def __init__(self, input_c) -> None:
super().__init__()
self.backbone = nn.Sequential(nn.Linear(in_features=input_c, out_features=32),
nn.LayerNorm(32),
nn.Linear(in_features=32, out_features=64),
nn.LayerNorm(64),
)
self.head = nn.Sequential(nn.Linear(in_features=64, out_features=1))
def forward(self, x):
x_feat = self.backbone(x)
x_feat = torch.max(x_feat, dim=-2)[0]
x_pred = self.head(x_feat)
return x_pred,x_feat
这个MModel类是一个使用PyTorch框架定义的神经网络模型。其主要功能和特点如下：
MModel类继承自torch.nn.Module，表示它是一个神经网络模型。
在初始化时，模型定义了两个主要部分：backbone和head。backbone是一个由两个线性层（nn.Linear）和两个层归一化层（nn.LayerNorm）组成的序列，用于提取输入数据的特征。
head部分是一个单独的线性层，用于将backbone输出的特征转换为最终的预测值。
在前向传播函数forward中，输入x首先通过backbone得到特征表示x_feat。
随后，x_feat在最后一个维度上取最大值（可能是为了降维或提取关键特征），然后传递给head以产生预测值x_pred。模型返回预测值x_pred和特征表示x_feat。
模型训练
data_dicts_train = gen_data(1000)
data_dicts_val = gen_data(100)
val_interval = 5
model = MModel(input_c=5).to(device)
optimizer = optim.SGD(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
epoch_max = 10
best_accruray = 0.0
for epoch in tqdm(range(epoch_max)):
loss_sum = 0.0
model.train()
for data_dict in (data_dicts_train):
input = data_dict['input']
output = data_dict['output']
input = torch.Tensor(input).unsqueeze(0).to(device)
output = torch.Tensor(output).unsqueeze(0).to(device)
pred,_ = model(input)
output_new = torch.sigmoid(output)
loss = criterion(pred, output_new)
loss_sum += loss.detach().item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'Epoch {epoch + 1}, Train Loss: {loss_sum / len(data_dicts_train)}')
if (epoch+1) % val_interval == 0:
loss_sum = 0.0
model.eval()
with torch.no_grad():
sum = 0
for data_dict in (data_dicts_val):
input = data_dict['input']
output = data_dict['output']
input = torch.Tensor(input).unsqueeze(0).to(device)
output = torch.Tensor(output).unsqueeze(0).to(device)
pred,_ = model(input)
output_new = torch.sigmoid(output)
loss = criterion(pred, output_new)
loss_sum += loss.detach().item()
if abs(pred.item()-float(output_new))<0.05:
sum+=1
accruray = float(sum)/len(data_dicts_val)
if accruray > best_accruray:
torch.save(model.state_dict(),'weights/best.pth')
print("save best.pth to dir:data")
best_accruray = accruray
print(f'Epoch {epoch + 1}, Val Loss: {loss_sum / len(data_dicts_val)}')
print(f'accruray is {accruray*100}%, best_accruray is {best_accruray*100}%')
这段代码实现了一个简单的神经网络训练过程，包括以下几个步骤：
生成训练集和验证集数据，其中每个数据项是一个包含输入和输出的字典。
创建一个MModel实例，并将其参数移至指定的设备上（如GPU）。
在多个训练周期（epoch）中，使用随机梯度下降（SGD）优化器和均方误差损失（MSELoss）来训练模型。
在每个训练周期结束时，计算训练损失，并在指定的验证间隔（val_interval）内评估模型在验证集上的性能。
如果模型在验证集上的性能有所提高，则保存其权重到文件’weights/best.pth’。
ONNX模型导出
在将PyTorch的.pth模型转换为TensorRT的.plan或.trt模型之前，先将其转换为ONNX（Open Neural Network Exchange）格式的主要目的是为了实现模型在不同深度学习框架之间的互操作性和可移植性。
ONNX是一个开源的模型表示格式，旨在使深度学习模型能够在不同框架之间无缝转换和使用。通过首先将PyTorch模型转换为ONNX格式，可以确保模型的结构和参数以一种通用的方式被保存下来，而不依赖于原始的PyTorch框架。
命令导出
trtexec --onnx=best.onnx --saveEngine=best.trt
源码导出
model = MModel(input_c=5).to(device)
model.load_state_dict(torch.load('weights/best.pth', map_location=device))
dummy = torch.zeros(1, 10, 5)
torch.onnx.export(
model, (dummy,), "weights/best.onnx",
input_names=["input"],
output_names=["output","feats"],
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}, "feats":{0:"batch"}},
opset_version=11
)
print('pth to onnx')
这段代码的作用是将已训练的PyTorch模型转换为ONNX格式。
加载预训练模型的权重（保存在’weights/best.pth’）。
使用torch.onnx.export函数将模型转换为ONNX文件（‘weights/best.onnx’），同时指定输入和输出名称，并允许输入维度（batch size）是动态的。
- 注意输入输出命名和模型的输入输出对应
- 动态shape通常只指定batch维度
tensorrt模型序列化
TensorRT模型序列化的作用主要是将训练好的深度学习模型（通常是以ONNX或其他格式保存的）转换为一个能够在TensorRT中运行的优化后的引擎（engine）。这个引擎是一个高度优化和特定于硬件的二进制文件，用于在NVIDIA GPU上进行高效的推理。
序列化过程包括将模型的信息和结构转换成可以存储或传输的格式，同时还会针对目标硬件进行一系列的优化，例如层融合、内存优化等，以最大程度地提高推理速度和效率。序列化后的引擎可以直接加载到TensorRT中进行推理，而无需再次进行模型的编译和优化。
此外，TensorRT还支持在序列化时保存支持动态batch的引擎，这使得模型能够适应不同大小的输入数据，而无需为每种可能的输入大小都单独编译一个引擎。这种灵活性使得TensorRT在部署深度学习模型时更加便捷和高效。
string root_path = "/files/code/Csdn-Source-Code-Implementation/tensorrt/step0-diymodel/weights/";
string onnx_file = root_path+"best.onnx";
string trt_file = root_path+"best.trt";
TRTLogger logger;
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);
nvonnxparser::IParser* parser = nvonnxparser::createParser(*network,logger);
if(!parser->parseFromFile(onnx_file.c_str(),1)){
printf("Failed to parse classifier.onnx\n");
// 注意这里的几个指针还没有释放，是有内存泄漏的，后面考虑更优雅的解决
return -1;
}
printf("succeed parse onnx\n");
config->setMaxWorkspaceSize(1<<28);
printf("WorkspaceSize = %.2f MB\n",(1<<28)/1024.0/1024.0);
nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
nvinfer1::ITensor* input = network->getInput(0);
nvinfer1::Dims input_dim = input->getDimensions();
input_dim.d[0] = 1;
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, input_dim);
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, input_dim);
input_dim.d[0] = 10;
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, input_dim);
config->addOptimizationProfile(profile);
print_ioinfo(network);
nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network,*config);
if(engine== nullptr){
printf("Build engine failed.\n");
return -1;
}
nvinfer1::IHostMemory* model_data = engine->serialize();
FILE *f = fopen(trt_file.c_str(), "wb");
fwrite(model_data->data(), 1, model_data->size(), f);
fclose(f);
model_data->destroy();
parser->destroy();
engine->destroy();
network->destroy();
config->destroy();
builder->destroy();
这段代码的主要作用是将一个ONNX（Open Neural Network Exchange）格式的深度学习模型转换为TensorRT的序列化引擎文件（.trt），以便于在NVIDIA GPU上进行高效的推理。以下是代码的详细解释：
- 设置路径和文件名：
定义了模型文件和序列化后引擎文件的存储路径。
onnx_file 是ONNX模型文件的路径。
trt_file 是序列化后TensorRT引擎文件的路径。 - 创建TensorRT组件：
创建了一个日志记录器（logger），用于在构建和解析模型时记录日志。
使用日志记录器创建了一个构建器（builder），它是TensorRT中用于创建网络定义、构建引擎等的主要组件。
创建了一个构建器配置（config），用于设置引擎构建时的各种参数。
创建了一个网络定义（network），用于定义模型的输入、输出和层。
创建了一个ONNX解析器（parser），用于将ONNX模型解析为TensorRT可以理解的格式。 - 解析ONNX模型：
使用ONNX解析器从文件加载并解析ONNX模型。
如果解析失败，则输出错误信息并返回-1。 - 配置工作空间和优化：
设置了构建引擎时所使用的最大工作空间大小（maxWorkspaceSize）。这个大小通常会影响引擎的性能和显存占用。
创建了一个优化配置文件（profile），并设置了输入张量的最小（kMIN）、最优（kOPT）和最大（kMAX）维度。这允许TensorRT为不同大小的输入数据优化引擎。 - 打印网络信息：
print_ioinfo(network); 用于打印网络的输入和输出信息。 - 构建TensorRT引擎：
使用构建器根据网络定义和配置参数构建TensorRT引擎。
如果构建失败，则输出错误信息并返回-1。 - 序列化引擎并保存到文件：
将构建好的引擎序列化为内存中的一块数据（model_data）。
将序列化后的数据写入到文件中（trt_file）。 - 释放资源：
释放了序列化数据、解析器、引擎、网络定义、配置和构建器所占用的内存，以防止内存泄漏。
这段代码的主要目的是将ONNX格式的深度学习模型转换为TensorRT引擎，以便于在NVIDIA GPU上进行高效的推理。同时，它也展示了TensorRT的一些基本用法，如创建组件、解析模型、配置参数、构建引擎和序列化引擎等。
Tensorrt推理
string root_path = "/files/code/Csdn-Source-Code-Implementation/tensorrt/step0-diymodel/weights/";
string trt_file = root_path + "best.trt";
TRTLogger logger;
std::vector<uint8_t> data = load_file(trt_file);
nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(data.data(), data.size());
if (engine == nullptr) {
printf("deserialize failed\n");
runtime->destroy();
return -1;
}
print_ioinfo(engine);
nvinfer1::Dims dim = engine->getBindingDimensions(0);
int length = dim.d[1];
int channel = dim.d[2];
int input_capacity = 1 * length * channel;
float *input_host = new float[input_capacity];
float *input_device = nullptr;
cudaMalloc(reinterpret_cast<void **>(&input_device), input_capacity * sizeof(float));
nvinfer1::Dims dim1 = engine->getBindingDimensions(1);
int feats_capacity = 1 * dim1.d[1];
vector<float> feats_host(feats_capacity, -1);
float *feats_device = nullptr;
cudaMalloc(reinterpret_cast<void **>(&feats_device), feats_capacity * sizeof(float));
nvinfer1::Dims dim2 = engine->getBindingDimensions(2);
int output_capacity = 1 * dim2.d[1];
float output_host[output_capacity];
float *output_device = nullptr;
cudaMalloc(reinterpret_cast<void **>(&output_device), output_capacity * sizeof(float));
vector<float *> binding{input_device, feats_device, output_device};
// input is 10 feats, 5 channel in one feat
float *output = new float(-100000);
...
generate input output
...
cudaStream_t stream = nullptr;
cudaStreamCreate(&stream);
cudaMemcpyAsync(input_device, input_host, input_capacity * sizeof(float), cudaMemcpyHostToDevice, stream);
nvinfer1::IExecutionContext *context = engine->createExecutionContext();
// 明确当前推理时，使用的数据输入大小
auto input_dims = context->getBindingDimensions(0);
input_dims.d[0] = 1;
// 设置当前推理时，input大小
context->setBindingDimensions(0, input_dims);
bool success = context->enqueueV2((void **) binding.data(), stream, nullptr);
cudaMemcpyAsync(output_host, output_device, output_capacity * sizeof(float), cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
*output_host = exp(*output_host);
float simularity_a = cos_similarity(output_host, output, output_capacity);
printf("simularity_a: %f\n",simularity_a);
这段代码展示了如何使用NVIDIA的TensorRT库进行深度学习模型的推理。以下是对代码的详细解释：
- 创建运行时和引擎
创建一个TensorRT运行时实例，其中logger用于记录日志。
使用序列化数据（可能是之前保存的引擎文件）来反序列化一个CUDA引擎实例。 - 获取模型绑定维度
获取输入、中间层和输出的维度信息。 - 分配主机和设备内存
input_host是在CPU上分配的内存，用于存储输入数据。
input_device是在GPU上分配的内存，用于存储模型推理时的输入数据。 - 初始化绑定向量
创建一个向量来存储指向输入、中间层和输出在GPU上内存的指针。 - 分配和初始化output变量（这里存在潜在问题）
这里只分配了一个float大小的内存，并用-100000进行了初始化。 - 生成输入和（可能）输出
…
generate input output
…
这部分代码填充input_host数组和输出的gt数据。 - 异步数据传输和推理
创建执行上下文、设置输入维度、执行推理
使用CUDA流进行异步的数据传输和推理。
cudaMemcpyAsync用于将数据从主机复制到设备（输入）和从设备复制回主机（输出）。
cudaStreamSynchronize确保在继续之前，所有在流上的操作都已完成。 - 后处理
对output_host元素应用指数函数，与python侧一致的后处理。
使用cos_similarity函数计算output_host和output之间的余弦相似度。
如需获取全套代码请参考
扫一扫
3243
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
