apollo::perception::inference::TorchNet类 参考

#include <torch_net.h>

类 apollo::perception::inference::TorchNet 继承关系图:

apollo::perception::inference::TorchNet 的协作图:

Public 成员函数
	TorchNet (const std::string &model_file, const std::vector< std::string > &outputs, const std::vector< std::string > &inputs)
virtual	~TorchNet ()
bool	Init (const std::map< std::string, std::vector< int > > &shapes) override
void	Infer () override
base::BlobPtr< float >	get_blob (const std::string &name) override
Public 成员函数 继承自 apollo::perception::inference::Inference
	Inference ()=default
virtual	~Inference ()=default
virtual void	SetStream (cudaStream_t stream)
void	set_max_batch_size (const int &batch_size)
void	set_gpu_id (const int &gpu_id)
void	set_model_info (const std::string &proto_file, const std::vector< std::string > &net_input_names, const std::vector< std::string > &net_output_names)

Protected 成员函数
bool	reshape ()
bool	shape (const std::string &name, std::vector< int > *res)

Protected 属性
torch::jit::script::Module	net_
Protected 属性 继承自 apollo::perception::inference::Inference
int	max_batch_size_ = 1
int	gpu_id_ = 0
std::string	proto_file_ = ""
std::vector< std::string >	net_output_names_
std::vector< std::string >	net_input_names_

详细描述

在文件 torch_net.h 第 33 行定义.

构造及析构函数说明

TorchNet()

apollo::perception::inference::TorchNet::TorchNet	(	const std::string &	model_file,
		const std::vector< std::string > &	outputs,
		const std::vector< std::string > &	inputs
	)

在文件 torch_net.cc 第 28 行定义.

    : model_file_(model_file), output_names_(outputs), input_names_(inputs) {}

~TorchNet()

virtual apollo::perception::inference::TorchNet::~TorchNet ( )

inlinevirtual

在文件 torch_net.h 第 39 行定义.

{}

成员函数说明

get_blob()

std::shared_ptr< Blob< float > > apollo::perception::inference::TorchNet::get_blob ( const std::string &  name )

overridevirtual

实现了 apollo::perception::inference::Inference.

在文件 torch_net.cc 第 65 行定义.

                                                                   {
  auto iter = blobs_.find(name);
  if (iter == blobs_.end()) {
    return nullptr;
  }
  return iter->second;
}

Infer()

void apollo::perception::inference::TorchNet::Infer ( )

overridevirtual

实现了 apollo::perception::inference::Inference.

在文件 torch_net.cc 第 85 行定义.

                     {
  torch::Device device(device_type_, device_id_);
  // Get input data from blob to torch_blob.
  std::vector<torch::jit::IValue> torch_inputs;
  for (const auto &name : input_names_) {
    auto blob = get_blob(name);
    if (blob != nullptr) {
      std::vector<int64_t> shape(blob->shape().begin(), blob->shape().end());
      torch::Tensor torch_blob = torch::from_blob(
          blob->data()->mutable_cpu_data(), shape, torch::kFloat32);
      torch_blob = torch_blob.to(device);
      torch_inputs.push_back(torch_blob);
    }
  }
  // If `out_blob->mutable_cpu_data()` is invoked outside,
  // HEAD will be set to CPU, and `out_blob->mutable_gpu_data()`
  // after `enqueue` will copy data from CPU to GPU,
  // which will overwrite the `inference` results.
  // `out_blob->gpu_data()` will set HEAD to SYNCED,
  // then no copy happends after `enqueue`.
  for (const auto &name : output_names_) {
    auto blob = get_blob(name);
    if (blob != nullptr) {
      blob->gpu_data();
    }
  }
 
  // Infer
  std::vector<torch::Tensor> output =
      net_.forward(torch_inputs).toTensorVector();
 
  // Fill output
  for (size_t i = 0; i < output_names_.size(); ++i) {
    auto blob = get_blob(output_names_[i]);
    if (blob != nullptr && i < output.size()) {
      std::vector<int64_t> output_size = output[i].sizes().vec();
      std::vector<int> shape(output_size.begin(), output_size.end());
      blob->Reshape(shape);
      blob->set_gpu_data(output[i].data_ptr<float>());
    }
  }
  emptyCache();
}

Init()

bool apollo::perception::inference::TorchNet::Init ( const std::map< std::string, std::vector< int > > &  shapes )

overridevirtual

实现了 apollo::perception::inference::Inference.

在文件 torch_net.cc 第 33 行定义.

                                                                   {
  if (gpu_id_ >= 0) {
    device_type_ = torch::kCUDA;
    device_id_ = gpu_id_;
  } else {
    device_type_ = torch::kCPU;
  }
 
  // Init net
  torch::Device device(device_type_, device_id_);
  net_ = torch::jit::load(model_file_, device);
  net_.eval();
 
  // add blobs
  for (const auto &name : input_names_) {
    auto iter = shapes.find(name);
    if (iter != shapes.end()) {
      auto blob = std::make_shared<Blob<float>>(iter->second);
      blobs_.emplace(name, blob);
    }
  }
 
  for (const auto &name : output_names_) {
    auto iter = shapes.find(name);
    if (iter != shapes.end()) {
      auto blob = std::make_shared<Blob<float>>(iter->second);
      blobs_.emplace(name, blob);
    }
  }
  return true;
}

reshape()

bool apollo::perception::inference::TorchNet::reshape ( )

protected

在文件 torch_net.cc 第 73 行定义.

{ return true; }

shape()

bool apollo::perception::inference::TorchNet::shape ( const std::string &  name, std::vector< int > *  res  )

protected

在文件 torch_net.cc 第 75 行定义.

                                                               {
  auto blob = get_blob(name);
  if (blob == nullptr) {
    return false;
  }
 
  *res = blob->shape();
  return true;
}

类成员变量说明

net_

torch::jit::script::Module apollo::perception::inference::TorchNet::net_

protected

在文件 torch_net.h 第 49 行定义.

---

该类的文档由以下文件生成:

• modules/perception/common/inference/libtorch/torch_net.h
• modules/perception/common/inference/libtorch/torch_net.cc