tensorrt 1.0.1 → 1.0.3
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/ext/tensorrt_rb/tensorrt_rb.cpp +32 -1
- data/lib/tensorrt/version.rb +1 -1
- metadata +1 -1
checksums.yaml
CHANGED
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
---
|
|
2
2
|
SHA256:
|
|
3
|
-
metadata.gz:
|
|
4
|
-
data.tar.gz:
|
|
3
|
+
metadata.gz: 0c6511e0f34d005d4054981835062177fbe44aa9bb58d3f13ab0a297de309836
|
|
4
|
+
data.tar.gz: 9eb97906e0cdc2bd9f89ff22c6da1d721080bc67fffc688aa14a722f9a1a3c52
|
|
5
5
|
SHA512:
|
|
6
|
-
metadata.gz:
|
|
7
|
-
data.tar.gz:
|
|
6
|
+
metadata.gz: 2d224a85ed3e90ff4645de1cd82e03e9bcd2f020eb181659d03f96eae058a3830bc2c5c63095ab1af7f5aa64efe3bddd697b3297dc91bf0179a7d82598e5af54
|
|
7
|
+
data.tar.gz: e68558aabfe741762facd9a0a5af14322e96a10928420d24a7ca9e23c083b150f8aa99d6d714599a53c55e1f5fbf5f774c7a55df4668f387ae60e9e776fd97b2
|
|
@@ -115,13 +115,43 @@ public:
|
|
|
115
115
|
|
|
116
116
|
size_t get_tensor_bytes(const std::string& name) const {
|
|
117
117
|
auto dims = engine->getTensorShape(name.c_str());
|
|
118
|
-
|
|
118
|
+
auto dtype = engine->getTensorDataType(name.c_str());
|
|
119
|
+
size_t element_size;
|
|
120
|
+
switch (dtype) {
|
|
121
|
+
case nvinfer1::DataType::kINT64: element_size = 8; break;
|
|
122
|
+
case nvinfer1::DataType::kFLOAT:
|
|
123
|
+
case nvinfer1::DataType::kINT32: element_size = 4; break;
|
|
124
|
+
case nvinfer1::DataType::kHALF:
|
|
125
|
+
case nvinfer1::DataType::kBF16: element_size = 2; break;
|
|
126
|
+
case nvinfer1::DataType::kINT8:
|
|
127
|
+
case nvinfer1::DataType::kBOOL:
|
|
128
|
+
case nvinfer1::DataType::kFP8:
|
|
129
|
+
case nvinfer1::DataType::kUINT8: element_size = 1; break;
|
|
130
|
+
default: element_size = 4; break;
|
|
131
|
+
}
|
|
132
|
+
size_t bytes = element_size;
|
|
119
133
|
for (int i = 0; i < dims.nbDims; i++) {
|
|
120
134
|
bytes *= dims.d[i];
|
|
121
135
|
}
|
|
122
136
|
return bytes;
|
|
123
137
|
}
|
|
124
138
|
|
|
139
|
+
std::string get_tensor_dtype(const std::string& name) const {
|
|
140
|
+
auto dtype = engine->getTensorDataType(name.c_str());
|
|
141
|
+
switch (dtype) {
|
|
142
|
+
case nvinfer1::DataType::kFLOAT: return "float32";
|
|
143
|
+
case nvinfer1::DataType::kHALF: return "float16";
|
|
144
|
+
case nvinfer1::DataType::kBF16: return "bfloat16";
|
|
145
|
+
case nvinfer1::DataType::kINT64: return "int64";
|
|
146
|
+
case nvinfer1::DataType::kINT32: return "int32";
|
|
147
|
+
case nvinfer1::DataType::kINT8: return "int8";
|
|
148
|
+
case nvinfer1::DataType::kUINT8: return "uint8";
|
|
149
|
+
case nvinfer1::DataType::kBOOL: return "bool";
|
|
150
|
+
case nvinfer1::DataType::kFP8: return "fp8";
|
|
151
|
+
default: return "unknown";
|
|
152
|
+
}
|
|
153
|
+
}
|
|
154
|
+
|
|
125
155
|
void set_tensor_address(const std::string& name, uint64_t ptr) {
|
|
126
156
|
void* addr = reinterpret_cast<void*>(ptr);
|
|
127
157
|
context->setTensorAddress(name.c_str(), addr);
|
|
@@ -172,6 +202,7 @@ extern "C" void Init_tensorrt_rb() {
|
|
|
172
202
|
.define_method("is_input?", &TRTEngine::is_input)
|
|
173
203
|
.define_method("get_tensor_shape", &TRTEngine::get_tensor_shape)
|
|
174
204
|
.define_method("get_tensor_bytes", &TRTEngine::get_tensor_bytes)
|
|
205
|
+
.define_method("get_tensor_dtype", &TRTEngine::get_tensor_dtype)
|
|
175
206
|
.define_method("set_tensor_address", &TRTEngine::set_tensor_address)
|
|
176
207
|
.define_method("execute", &TRTEngine::execute)
|
|
177
208
|
.define_method("enqueue", &TRTEngine::enqueue)
|
data/lib/tensorrt/version.rb
CHANGED