//根據配置的參數RTNet建立RTNet網絡
Inference *CreateInferenceByName(const std::string &name,
const std::string &proto_file,
const std::string &weight_file,
const std::vector<std::string> &outputs,
const std::vector<std::string> &inputs,
const std::string &model_root) {
if (name == "RTNet") {
return new RTNet(proto_file, weight_file, outputs, inputs);
} else if (name == "RTNetInt8") {
return new RTNet(proto_file, weight_file, outputs, inputs, model_root);
} else if (name == "TorchDet") {
return new TorchDet(proto_file, weight_file, outputs, inputs);
} else if (name == "TorchNet") {
return new TorchNet(proto_file, weight_file, outputs, inputs);
} else if (name == "Obstacle") {
return new ObstacleDetector(proto_file, weight_file, outputs, inputs);
}
return nullptr;
}
//構造函數加載權值和網絡參數
RTNet::RTNet(const std::string &net_file, const std::string &model_file,
const std::vector<std::string> &outputs,
const std::vector<std::string> &inputs)
: output_names_(outputs), input_names_(inputs) {
loadWeights(model_file, &weight_map_);
net_param_.reset(new NetParameter);
loadNetParams(net_file, net_param_.get());
}
//初始化網絡,通過parse_with_api從檔案中讀取網絡的模型并生成網絡
bool RTNet::Init(const std::map<std::string, std::vector<int>> &shapes) {
if (gpu_id_ < 0) {
AINFO << "must use gpu mode";
return false;
}
BASE_CUDA_CHECK(cudaSetDevice(gpu_id_));
// stream will only be destoried for gpu_id_ >= 0
cudaStreamCreate(&stream_);
builder_ = nvinfer1::createInferBuilder(rt_gLogger);
network_ = builder_->createNetwork();
parse_with_api(shapes);
builder_->setMaxBatchSize(max_batch_size_);
workspaceSize_ = 1 << 30;
builder_->setMaxWorkspaceSize(workspaceSize_);
cudaDeviceProp prop;
cudaGetDeviceProperties(&prop, gpu_id_);
bool int8_mode = checkInt8(prop.name, calibrator_);
builder_->setInt8Mode(int8_mode);
builder_->setInt8Calibrator(calibrator_);
builder_->setDebugSync(true);
nvinfer1::ICudaEngine *engine = builder_->buildCudaEngine(*network_);
context_ = engine->createExecutionContext();
buffers_.resize(input_names_.size() + output_names_.size());
init_blob(&input_names_);
init_blob(&output_names_);
return true;
}
void RTNet::parse_with_api(
const std::map<std::string, std::vector<int>> &shapes) {
CHECK_GT(net_param_->layer_size(), 0);
std::vector<LayerParameter> order;
TensorMap tensor_map;
TensorDimsMap tensor_dims_map;
ParseNetParam(*net_param_, &tensor_dims_map, &tensor_modify_map_, &order);
//添加網絡輸入層,輸入層的名字為data,形狀為[1,3,480,640]
addInput(tensor_dims_map, shapes, &tensor_map);
for (auto layer_param : order) {
std::vector<nvinfer1::ITensor *> inputs;
for (int j = 0; j < layer_param.bottom_size(); j++) {
CHECK_NOTNULL(tensor_map[tensor_modify_map_[layer_param.bottom(j)]]);
inputs.push_back(tensor_map[tensor_modify_map_[layer_param.bottom(j)]]);
}
//依次添加其他網絡層
addLayer(layer_param, inputs.data(), layer_param.bottom_size(),
&weight_map_, network_, &tensor_map, &tensor_modify_map_);
}
//輸出層的形狀為[1,13,480,640],每個通道表達所屬每個車道的資訊,圖檔上的單元格表示目前車道
//在目前位置是車道線的機率
CHECK_NE(output_names_.size(), static_cast<size_t>(0));
std::sort(output_names_.begin(), output_names_.end());
auto last = std::unique(output_names_.begin(), output_names_.end());
output_names_.erase(last, output_names_.end());
for (auto iter = output_names_.begin(); iter != output_names_.end();) {
if (tensor_map.find(tensor_modify_map_[*iter]) != tensor_map.end()) {
network_->markOutput(*tensor_map[tensor_modify_map_[*iter]]);
iter++;
} else {
AINFO << "Erase output: " << *iter;
iter = output_names_.erase(iter);
}
}
}
void RTNet::addLayer(const LayerParameter &layer_param,
nvinfer1::ITensor *const *inputs, int nbInputs,
WeightMap *weight_map, nvinfer1::INetworkDefinition *net,
TensorMap *tensor_map,
TensorModifyMap *tensor_modify_map) {
//卷積層
if (layer_param.type() == "Convolution") {
addConvLayer(layer_param, inputs, weight_map, net, tensor_map,
tensor_modify_map);
}
//反卷積層
else if (layer_param.type() == "Deconvolution") {
addDeconvLayer(layer_param, inputs, weight_map, net, tensor_map,
tensor_modify_map);
}
//激活層
else if (layer_param.type() == "Sigmoid" || layer_param.type() == "TanH" ||
layer_param.type() == "ReLU") {
addActiveLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "Concat") {
addConcatLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "Slice") {
addSliceLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "Reshape") {
addReshapeLayer(layer_param, inputs, net, tensor_map, tensor_modify_map);
} else if (layer_param.type() == "Padding") {
addPaddingLayer(layer_param, inputs, net, tensor_map, tensor_modify_map);
} else if (layer_param.type() == "Pooling") {
addPoolingLayer(layer_param, inputs, net, tensor_map, tensor_modify_map);
} else if (layer_param.type() == "Permute") {
addPermuteLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "InnerProduct") {
addInnerproductLayer(layer_param, inputs, weight_map, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "BatchNorm") {
addBatchnormLayer(layer_param, inputs, weight_map, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "Scale") {
addScaleLayer(layer_param, inputs, weight_map, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "Softmax") {
addSoftmaxLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "Eltwise") {
addEltwiseLayer(layer_param, inputs, net, tensor_map, tensor_modify_map);
} else if (layer_param.type() == "ArgMax") {
addArgmaxLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "Dropout") {
AINFO << "skip dropout";
} else if (layer_param.type() == "Power") {
addScaleLayer(layer_param, inputs, weight_map, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "DFMBPSROIAlign") {
addDFMBPSROIAlignLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "RCNNProposal") {
addRCNNProposalLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else if (layer_param.type() == "RPNProposalSSD") {
addRPNProposalSSDLayer(layer_param, inputs, nbInputs, net, tensor_map,
tensor_modify_map);
} else {
AWARN << "unknown layer type:" << layer_param.type();
}
} |