Skip to content

Commit 4d0e3fd

Browse files
author
Wish
committed
update
1 parent a0cebf2 commit 4d0e3fd

File tree

6 files changed

+30
-39
lines changed

6 files changed

+30
-39
lines changed

‎src/application/app_python/interface.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -700,9 +700,9 @@ PYBIND11_MODULE(libtrtpyc, m) {
700700
.def("resize", [](TRT::Tensor& self, const vector<int>& dims) {self.resize(dims);})
701701
.def("resize_single_dim", [](TRT::Tensor& self, int dim, int size) {self.resize_single_dim(dim, size);})
702702
.def("count", [](TRT::Tensor& self, int start_axis) {return self.count(start_axis);})
703-
.def("offset", [](TRT::Tensor& self, const vector<int>& indexs){return self.offset(indexs);})
704-
.def("cpu_at", [](TRT::Tensor& self, const vector<int>& indexs){return ptr_wrapper<float, ptr_base::host >(self.cpu<float>() + self.offset(indexs));})
705-
.def("gpu_at", [](TRT::Tensor& self, const vector<int>& indexs){return ptr_wrapper<float, ptr_base::device>(self.gpu<float>() + self.offset(indexs));})
703+
.def("offset", [](TRT::Tensor& self, const vector<int>& indexs){return self.offset_array(indexs);})
704+
.def("cpu_at", [](TRT::Tensor& self, const vector<int>& indexs){return ptr_wrapper<float, ptr_base::host >(self.cpu<float>() + self.offset_array(indexs));})
705+
.def("gpu_at", [](TRT::Tensor& self, const vector<int>& indexs){return ptr_wrapper<float, ptr_base::device>(self.gpu<float>() + self.offset_array(indexs));})
706706
.def_property_readonly("cpu", [](TRT::Tensor& self){return ptr_wrapper<float, ptr_base::host >(self.cpu<float>());})
707707
.def_property_readonly("gpu", [](TRT::Tensor& self){return ptr_wrapper<float, ptr_base::device>(self.gpu<float>());})
708708
.def_property_readonly("head", [](TRT::Tensor& self){return self.head();})

‎src/application/app_yolo.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ static void test_plugin(TRT::Mode mode){
117117
auto input1 = engine->input(1);
118118
auto output = engine->output(0);
119119

120+
INFO("offset %d", output->offset(1, 0));
120121
INFO("input0: %s", input0->shape_string());
121122
INFO("input1: %s", input1->shape_string());
122123
INFO("output: %s", output->shape_string());

‎src/tensorRT/common/trt_tensor.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -454,7 +454,7 @@ namespace TRT{
454454
return *this;
455455
}
456456

457-
int Tensor::offset(size_t size, const int* index_array){
457+
int Tensor::offset_array(size_t size, const int* index_array){
458458

459459
Assert(size <= shape_.size());
460460
int value = 0;
@@ -469,8 +469,8 @@ namespace TRT{
469469
return value;
470470
}
471471

472-
int Tensor::offset(const std::vector<int>& index_array){
473-
return offset(index_array.size(), index_array.data());
472+
int Tensor::offset_array(const std::vector<int>& index_array){
473+
return offset_array(index_array.size(), index_array.data());
474474
}
475475

476476
Tensor& Tensor::set_norm_mat(int n, const cv::Mat& image, float mean[3], float std[3]) {

‎src/tensorRT/common/trt_tensor.hpp‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,15 @@ namespace TRT {
101101
Tensor& set_to(float value);
102102
bool empty();
103103

104+
template<typename ... _Args>
105+
int offset(int index, _Args ... index_args){
106+
const int index_array[] = {index, index_args...};
107+
return offset_array(sizeof...(index_args) + 1, index_array);
108+
}
109+
110+
int offset_array(const std::vector<int>& index);
111+
int offset_array(size_t size, const int* index_array);
112+
104113
template<typename ... _Args>
105114
Tensor& resize(int dim_size, _Args ... dim_size_args){
106115
const int dim_size_array[] = {dim_size, dim_size_args...};
@@ -114,20 +123,11 @@ namespace TRT {
114123

115124
Tensor& to_gpu(bool copyedIfCPU = true);
116125
Tensor& to_cpu(bool copyedIfGPU = true);
117-
Tensor& to_half();
118126

127+
Tensor& to_half();
119128
Tensor& to_float();
120129
inline void* cpu() const { ((Tensor*)this)->to_cpu(); return data_->cpu(); }
121130
inline void* gpu() const { ((Tensor*)this)->to_gpu(); return data_->gpu(); }
122-
123-
template<typename ... _Args>
124-
int offset(int index, _Args ... index_args){
125-
const int index_array[] = {index, index_args...};
126-
return offset(sizeof...(index_args) + 1, index_array);
127-
}
128-
129-
int offset(const std::vector<int>& index);
130-
int offset(size_t size, const int* index_array);
131131

132132
template<typename DType> inline const DType* cpu() const { return (DType*)cpu(); }
133133
template<typename DType> inline DType* cpu() { return (DType*)cpu(); }

‎src/tensorRT/onnxplugin/onnxplugin.cpp‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,21 +13,25 @@ namespace ONNXPlugin {
1313
this->dtType_ = TRT::DataType::Float;
1414
}
1515

16-
int GTensor::offset(const std::vector<int>& index){
16+
int GTensor::offset_array(size_t size, const int* index_array){
1717

18-
Assert(index.size() <= shape_.size());
18+
Assert(size <= shape_.size());
1919
int value = 0;
2020
for(int i = 0; i < shape_.size(); ++i){
2121

22-
if(i < index.size())
23-
value += index[i];
22+
if(i < size)
23+
value += index_array[i];
2424

2525
if(i + 1 < shape_.size())
2626
value *= shape_[i+1];
2727
}
2828
return value;
2929
}
3030

31+
int GTensor::offset_array(const std::vector<int>& index){
32+
return offset_array(index.size(), index.data());
33+
}
34+
3135
GTensor::GTensor(TRT::float16* ptr, int ndims, int* dims) {
3236
this->ptr_ = ptr;
3337
this->shape_.insert(shape_.end(), dims, dims + ndims);

‎src/tensorRT/onnxplugin/onnxplugin.hpp‎

Lines changed: 5 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,13 @@ namespace ONNXPlugin {
3030
int count(int start_axis = 0) const;
3131

3232
template<typename ... _Args>
33-
int offset(int t, _Args&& ... args){
34-
offset_index_.clear();
35-
return offsetimpl(t, args...);
33+
int offset(int index, _Args&& ... index_args){
34+
const int index_array[] = {index, index_args...};
35+
return offset_array(sizeof...(index_args) + 1, index_array);
3636
}
3737

38-
int offset(const std::vector<int>& index);
38+
int offset_array(const std::vector<int>& index);
39+
int offset_array(size_t size, const int* index_array);
3940

4041
template<typename _T>
4142
inline _T* ptr() const { return (_T*)ptr_; }
@@ -56,21 +57,6 @@ namespace ONNXPlugin {
5657
void* ptr_ = nullptr;
5758
TRT::DataType dtType_ = TRT::DataType::Float;
5859
std::vector<int> shape_;
59-
60-
private:
61-
std::vector<int> offset_index_;
62-
63-
private:
64-
int offsetimpl(int value){
65-
offset_index_.push_back(value);
66-
return offset(offset_index_);
67-
}
68-
69-
template<typename ... _Args>
70-
int offsetimpl(int t, _Args&& ... args){
71-
offset_index_.push_back(t);
72-
return offsetimpl(args...);
73-
}
7460
};
7561

7662
struct LayerConfig {

0 commit comments

Comments
 (0)