tensorrt 1.0.2 → 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 +17 -0
- 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
|
|
@@ -136,6 +136,22 @@ public:
|
|
|
136
136
|
return bytes;
|
|
137
137
|
}
|
|
138
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
|
+
|
|
139
155
|
void set_tensor_address(const std::string& name, uint64_t ptr) {
|
|
140
156
|
void* addr = reinterpret_cast<void*>(ptr);
|
|
141
157
|
context->setTensorAddress(name.c_str(), addr);
|
|
@@ -186,6 +202,7 @@ extern "C" void Init_tensorrt_rb() {
|
|
|
186
202
|
.define_method("is_input?", &TRTEngine::is_input)
|
|
187
203
|
.define_method("get_tensor_shape", &TRTEngine::get_tensor_shape)
|
|
188
204
|
.define_method("get_tensor_bytes", &TRTEngine::get_tensor_bytes)
|
|
205
|
+
.define_method("get_tensor_dtype", &TRTEngine::get_tensor_dtype)
|
|
189
206
|
.define_method("set_tensor_address", &TRTEngine::set_tensor_address)
|
|
190
207
|
.define_method("execute", &TRTEngine::execute)
|
|
191
208
|
.define_method("enqueue", &TRTEngine::enqueue)
|
data/lib/tensorrt/version.rb
CHANGED