HyperPose/src/uff_runner.cpp at 1.0.0 · tensorlayer/HyperPose · GitHub

History

240 lines (210 loc) · 7.46 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

#include <algorithm>

#include <cassert>

#include <cstdio>

#include <cstdlib>

#include <fstream>

#include <iostream>

#include <memory>

#include <string>

#include <vector>

#include <ttl/cuda_tensor>

#include <ttl/experimental/copy>

#include <ttl/range>

#include <NvInfer.h>

#include <NvUffParser.h>

#include <NvUtils.h>

#include <openpose-plus.h>

#include "logger.h"

#include "trace.hpp"

using input_info_t = std::vector<std::pair<std::string, std::vector<int>>>;

Logger gLogger;

inline int64_t volume(const nvinfer1::Dims &d)

{

int64_t v = 1;

for (int i = 0; i < d.nbDims; i++) { v *= d.d[i]; }

return v;

}

inline size_t elementSize(nvinfer1::DataType t)

{

switch (t) {

// TODO: check nvinfer1 version

// case nvinfer1::DataType::kINT32:

// return 4;

case nvinfer1::DataType::kFLOAT:

return 4;

case nvinfer1::DataType::kHALF:

return 2;

case nvinfer1::DataType::kINT8:

return 1;

}

assert(0);

return 0;

}

std::string to_string(const nvinfer1::Dims &d)

{

std::string s{"("};

if (d.nbDims != 0) {

for (int64_t i = 0; i < d.nbDims; i++)

(s += std::to_string(d.d[i])) += ", ";

s.pop_back();

}

return s + ")";

}

std::string to_string(const nvinfer1::DataType dtype)

{

return std::to_string(int(dtype));

}

template <typename T> struct destroy_deleter {

void operator()(T *ptr) { ptr->destroy(); }

};

template <typename T>

using destroy_ptr = std::unique_ptr<T, destroy_deleter<T>>;

nvinfer1::ICudaEngine *loadModelAndCreateEngine(const char *uffFile,

int max_batch_size,

nvuffparser::IUffParser *parser,

bool use_f16)

{

destroy_ptr<nvinfer1::IBuilder> builder(

nvinfer1::createInferBuilder(gLogger));

destroy_ptr<nvinfer1::INetworkDefinition> network(builder->createNetwork());

if (use_f16) {

if (!parser->parse(uffFile, *network, nvinfer1::DataType::kHALF)) {

return nullptr;

}

// builder->setFp16Mode(true);

builder->setHalf2Mode(true); // For older version of tensorRT

} else {

if (!parser->parse(uffFile, *network, nvinfer1::DataType::kFLOAT)) {

return nullptr;

}

builder->setMaxBatchSize(max_batch_size);

return builder->buildCudaEngine(*network);

}

nvinfer1::ICudaEngine *

create_engine(const std::string &model_file, const input_info_t &input_info,

const std::vector<std::string> &output_names, int max_batch_size,

bool use_f16)

{

TRACE_SCOPE(__func__);

destroy_ptr<nvuffparser::IUffParser> parser(nvuffparser::createUffParser());

for (const auto &info : input_info) {

const auto dims = info.second;

parser->registerInput(

info.first.c_str(),

// Always provide your dimensions in CHW even if your

// network input was in HWC in yout original framework.

nvinfer1::DimsCHW(dims[0], dims[1], dims[2]),

nvuffparser::UffInputOrder::kNCHW //

);

}

for (auto &name : output_names) { parser->registerOutput(name.c_str()); }

auto engine = loadModelAndCreateEngine(model_file.c_str(), max_batch_size,

parser.get(), use_f16);

if (!engine) {

gLogger.log(nvinfer1::ILogger::Severity::kERROR,

"failed to created engine");

exit(1);

}

return engine;

}

class uff_runner_impl : public pose_detection_runner

{

public:

uff_runner_impl(const std::string &model_file,

const input_info_t &input_info,

const std::vector<std::string> &output_names,

int max_batch_size, bool use_f16);

~uff_runner_impl() override;

void operator()(const std::vector<void *> &inputs,

const std::vector<void *> &outputs,

int batch_size) override;

private:

const int max_batch_size;

destroy_ptr<nvinfer1::ICudaEngine> engine_;

using cuda_buffer_t = ttl::cuda_tensor<char, 2>; // [batch_size, data_size]

std::vector<cuda_buffer_t> buffers_;

void createBuffers_(int batch_size);

};

uff_runner_impl::uff_runner_impl(const std::string &model_file,

const input_info_t &input_info,

const std::vector<std::string> &output_names,

int max_batch_size, bool use_f16)

: max_batch_size(max_batch_size),

engine_(create_engine(model_file, input_info, output_names,

max_batch_size, use_f16))

{

createBuffers_(max_batch_size);

}

uff_runner_impl::~uff_runner_impl() { nvuffparser::shutdownProtobufLibrary(); }

void uff_runner_impl::createBuffers_(int batch_size)

{

TRACE_SCOPE(__func__);

for (auto i : ttl::range(engine_->getNbBindings())) {

const nvinfer1::Dims dims = engine_->getBindingDimensions(i);

const nvinfer1::DataType dtype = engine_->getBindingDataType(i);

const std::string name(engine_->getBindingName(i));

std::cout << "binding " << i << ":"

<< " name: " << name << " type" << to_string(dtype)

<< to_string(dims) << std::endl;

buffers_.emplace_back(batch_size, volume(dims) * elementSize(dtype));

}

void uff_runner_impl::operator()(const std::vector<void *> &inputs,

const std::vector<void *> &outputs,

int batch_size)

{

TRACE_SCOPE("uff_runner_impl::operator()");

assert(batch_size <= max_batch_size);

{

TRACE_SCOPE("copy input from host");

int idx = 0;

for (auto i : ttl::range(buffers_.size())) {

if (engine_->bindingIsInput(i)) {

const auto buffer = buffers_[i].slice(0, batch_size);

ttl::tensor_view<char, 2> input(

reinterpret_cast<char *>(inputs[idx++]), buffer.shape());

ttl::copy(buffer, input);

}

{

TRACE_SCOPE("uff_runner_impl::context->execute");

auto context = engine_->createExecutionContext();

std::vector<void *> buffer_ptrs_(buffers_.size());

std::transform(buffers_.begin(), buffers_.end(), buffer_ptrs_.begin(),

[](const auto &b) { return b.data(); });

context->execute(batch_size, buffer_ptrs_.data());

context->destroy();

}

{

TRACE_SCOPE("copy output to host");

int idx = 0;

for (auto i : ttl::range(buffers_.size())) {

if (!engine_->bindingIsInput(i)) {

const auto buffer = buffers_[i].slice(0, batch_size);

ttl::tensor_ref<char, 2> output(

reinterpret_cast<char *>(outputs[idx++]), buffer.shape());

ttl::copy(output, ttl::view(buffer));

}

pose_detection_runner *

create_pose_detection_runner(const std::string &model_file, int input_height,

int input_width, int max_batch_size, bool use_f16)

{

const input_info_t input_info = {

{

"image",

{3, input_height, input_width} /* must be (C, H, W) */,

};

const std::vector<std::string> output_names = {

"outputs/conf",

"outputs/paf",

};

return new uff_runner_impl(model_file, input_info, output_names,

max_batch_size, use_f16);

}

HyperPose/src/uff_runner.cpp at 1.0.0 · tensorlayer/HyperPose · GitHub

Uh oh!

Expand file tree

Breadcrumbs

uff_runner.cpp

Latest commit

History

Breadcrumbs

uff_runner.cpp

File metadata and controls