Skip to content

Commit aa01024

Browse files
nihuilovedream-mslancerstadiumAtomAlpaca
authored
pnnx: support npy input tensors (#6700)
Co-authored-by: lovedream-ms <209071654+lovedream-ms@users.noreply.github.com> Co-authored-by: lancerstadium <57707690+lancerstadium@users.noreply.github.com> Co-authored-by: AtomAlpaca <66774326+AtomAlpaca@users.noreply.github.com>
1 parent 0d29a8d commit aa01024

12 files changed

Lines changed: 1172 additions & 28 deletions

‎tools/pnnx/src/ir.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2771,14 +2771,14 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
27712771
fprintf(pyfp, " torch.manual_seed(0)\n");
27722772

27732773
int input_shapes_i = 0;
2774-
27752774
std::vector<std::string> input_names;
27762775
for (const Operator* op : ops)
27772776
{
27782777
if (op->type != "pnnx.Input")
27792778
continue;
27802779

27812780
const Operand* r = op->outputs[0];
2781+
std::string input_name = std::string("v_") + sanitize_identifier(r->name);
27822782

27832783
std::vector<int> input_shape;
27842784
if (input_shapes.empty())
@@ -2794,7 +2794,6 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
27942794
}
27952795
}
27962796

2797-
std::string input_name = std::string("v_") + sanitize_identifier(r->name);
27982797
if (type_is_integer(r->type))
27992798
{
28002799
fprintf(pyfp, " %s = torch.randint(10, (", input_name.c_str());

‎tools/pnnx/src/load_tnn.cpp‎

Lines changed: 63 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,26 @@ static int get_tnn_tensor_type(int dt)
141141
return 0; // unknown type
142142
}
143143

144+
static int input_type_to_pnnx_type(const std::string& t)
145+
{
146+
if (t == "f32") return 1;
147+
if (t == "f64") return 2;
148+
if (t == "f16") return 3;
149+
if (t == "i32") return 4;
150+
if (t == "i64") return 5;
151+
if (t == "i16") return 6;
152+
if (t == "i8") return 7;
153+
if (t == "u8") return 8;
154+
if (t == "bool") return 9;
155+
if (t == "c64") return 10;
156+
if (t == "c128") return 11;
157+
if (t == "c32") return 12;
158+
if (t == "bf16") return 13;
159+
160+
fprintf(stderr, "unsupported input type %s\n", t.c_str());
161+
return 0;
162+
}
163+
144164
Attribute::Attribute(FILE* bp)
145165
{
146166
unsigned int magic;
@@ -178,7 +198,9 @@ Attribute::Attribute(FILE* bp)
178198
fread((void*)data.data(), 1, length, bp);
179199
}
180200

181-
int load_tnn(const std::string& tnnpath, Graph& pnnx_graph)
201+
int load_tnn(const std::string& tnnpath, Graph& pnnx_graph,
202+
const std::vector<std::vector<int64_t> >& input_shapes,
203+
const std::vector<std::string>& input_types)
182204
{
183205
fprintf(stderr, "############# pass_level0 tnn\n");
184206

@@ -244,14 +266,53 @@ int load_tnn(const std::string& tnnpath, Graph& pnnx_graph)
244266
int datatype = 0;
245267
sscanf(pline, "%d%n", &datatype, &ncomsumed);
246268

269+
int tensor_type = get_tnn_tensor_type(datatype);
270+
271+
if (!input_shapes.empty())
272+
{
273+
if (input_shapes.size() != 1)
274+
{
275+
fprintf(stderr, "tnn input expect 1 tensor but got %d\n", (int)input_shapes.size());
276+
return -1;
277+
}
278+
if (input_types.size() != input_shapes.size())
279+
{
280+
fprintf(stderr, "tnn input type count mismatch, expect %d but got %d\n", (int)input_shapes.size(), (int)input_types.size());
281+
return -1;
282+
}
283+
284+
shape.clear();
285+
const std::vector<int64_t>& input_shape = input_shapes[0];
286+
for (size_t i = 0; i < input_shape.size(); i++)
287+
{
288+
if (input_shape[i] < 0 || input_shape[i] > INT_MAX)
289+
{
290+
fprintf(stderr, "invalid tnn input shape dimension %ld\n", input_shape[i]);
291+
return -1;
292+
}
293+
shape.push_back((int)input_shape[i]);
294+
}
295+
296+
const int input_tensor_type = input_type_to_pnnx_type(input_types[0]);
297+
if (input_tensor_type == 0)
298+
return -1;
299+
if (tensor_type != 0 && tensor_type != input_tensor_type)
300+
{
301+
fprintf(stderr, "tnn input type mismatch, model type %d but got %s\n", tensor_type, input_types[0].c_str());
302+
return -1;
303+
}
304+
305+
tensor_type = input_tensor_type;
306+
}
307+
247308
Operator* op = pnnx_graph.new_operator("pnnx.Input", "input0");
248309

249310
Operand* r = pnnx_graph.new_operand(blob_name);
250311

251312
r->producer = op;
252313

253314
r->shape = shape;
254-
r->type = get_tnn_tensor_type(datatype);
315+
r->type = tensor_type;
255316

256317
op->outputs.push_back(r);
257318
}

‎tools/pnnx/src/load_tnn.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,9 @@
88

99
namespace pnnx {
1010

11-
int load_tnn(const std::string& tnnpath, Graph& g);
11+
int load_tnn(const std::string& tnnpath, Graph& g,
12+
const std::vector<std::vector<int64_t> >& input_shapes,
13+
const std::vector<std::string>& input_types);
1214

1315
} // namespace pnnx
1416

‎tools/pnnx/src/load_torchscript.cpp‎

Lines changed: 149 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
#include <torch/script.h>
1313
#include <torch/csrc/api/include/torch/version.h>
1414
#include <torch/csrc/jit/serialization/import_read.h>
15+
16+
#include <limits>
1517
#ifdef PNNX_TORCHVISION
1618
namespace vision {
1719
int64_t cuda_version();
@@ -418,11 +420,48 @@ static c10::ScalarType input_type_to_c10_ScalarType(const std::string& t)
418420
if (t == "i64") return torch::kInt64;
419421
if (t == "i8") return torch::kInt8;
420422
if (t == "u8") return torch::kUInt8;
423+
if (t == "bool") return torch::kBool;
421424

422425
fprintf(stderr, "unsupported type %s fallback to f32\n", t.c_str());
423426
return torch::kFloat32;
424427
}
425428

429+
static size_t input_type_to_elemsize(const std::string& t)
430+
{
431+
if (t == "f32") return 4;
432+
if (t == "f64") return 8;
433+
if (t == "f16") return 2;
434+
if (t == "i32") return 4;
435+
if (t == "i64") return 8;
436+
if (t == "i16") return 2;
437+
if (t == "i8") return 1;
438+
if (t == "u8") return 1;
439+
if (t == "bool") return 1;
440+
if (t == "c64") return 8;
441+
if (t == "c128") return 16;
442+
if (t == "c32") return 4;
443+
if (t == "bf16") return 2;
444+
return 0;
445+
}
446+
447+
static bool shape_element_count(const std::vector<int64_t>& shape, size_t& elem_count)
448+
{
449+
elem_count = 1;
450+
for (size_t i = 0; i < shape.size(); i++)
451+
{
452+
if (shape[i] < 0)
453+
return false;
454+
455+
const size_t dim = (size_t)shape[i];
456+
if (dim != 0 && elem_count > std::numeric_limits<size_t>::max() / dim)
457+
return false;
458+
459+
elem_count *= dim;
460+
}
461+
462+
return true;
463+
}
464+
426465
static const char* get_at_tensor_type_str(const at::ScalarType& st)
427466
{
428467
if (st == c10::ScalarType::Float) return "f32";
@@ -437,6 +476,7 @@ static const char* get_at_tensor_type_str(const at::ScalarType& st)
437476
if (st == c10::ScalarType::ComplexDouble) return "c128";
438477
if (st == c10::ScalarType::ComplexHalf) return "c32";
439478
if (st == c10::ScalarType::BFloat16) return "bf16";
479+
if (st == c10::ScalarType::Bool) return "bool";
440480

441481
// unknown
442482
fprintf(stderr, "unsupported tensor elem data type %d\n", (int)st);
@@ -583,8 +623,10 @@ int load_torchscript(const std::string& ptpath, Graph& pnnx_graph,
583623
const std::string& device,
584624
const std::vector<std::vector<int64_t> >& input_shapes,
585625
const std::vector<std::string>& input_types,
626+
const std::vector<std::vector<char> >& input_contents,
586627
const std::vector<std::vector<int64_t> >& input_shapes2,
587628
const std::vector<std::string>& input_types2,
629+
const std::vector<std::vector<char> >& input_contents2,
588630
const std::vector<std::string>& customop_modules,
589631
const std::vector<std::string>& module_operators,
590632
const std::string& foldable_constants_zippath,
@@ -646,29 +688,122 @@ int load_torchscript(const std::string& ptpath, Graph& pnnx_graph,
646688
}
647689

648690
std::vector<at::Tensor> input_tensors;
649-
for (size_t i = 0; i < traced_input_shapes.size(); i++)
691+
if (!input_contents.empty())
692+
{
693+
if (input_contents.size() != traced_input_shapes.size())
694+
{
695+
fprintf(stderr, "input expect %d tensors but got %d\n", (int)traced_input_shapes.size(), (int)input_contents.size());
696+
return -1;
697+
}
698+
699+
for (size_t i = 0; i < traced_input_shapes.size(); i++)
700+
{
701+
const std::vector<int64_t>& shape = traced_input_shapes[i];
702+
const std::string& type = traced_input_types[i];
703+
704+
size_t elem_count = 0;
705+
if (!shape_element_count(shape, elem_count))
706+
{
707+
fprintf(stderr, "invalid input shape for tensor %d\n", (int)i);
708+
return -1;
709+
}
710+
711+
const size_t elemsize = input_type_to_elemsize(type);
712+
if (elemsize == 0)
713+
{
714+
fprintf(stderr, "unsupported input type %s\n", type.c_str());
715+
return -1;
716+
}
717+
718+
if (elem_count != input_contents[i].size() / elemsize || elem_count * elemsize != input_contents[i].size())
719+
{
720+
fprintf(stderr, "input tensor %d data size mismatch for shape [", (int)i);
721+
for (size_t j = 0; j < shape.size(); j++)
722+
{
723+
fprintf(stderr, "%ld", shape[j]);
724+
if (j + 1 != shape.size())
725+
fprintf(stderr, ",");
726+
}
727+
fprintf(stderr, "]%s\n", type.c_str());
728+
return -1;
729+
}
730+
731+
at::Tensor t = torch::from_blob((void*)input_contents[i].data(), shape, torch::TensorOptions().dtype(input_type_to_c10_ScalarType(type))).clone();
732+
if (device == "gpu")
733+
t = t.cuda();
734+
735+
input_tensors.push_back(t);
736+
}
737+
}
738+
else
650739
{
651-
const std::vector<int64_t>& shape = traced_input_shapes[i];
652-
const std::string& type = traced_input_types[i];
740+
for (size_t i = 0; i < traced_input_shapes.size(); i++)
741+
{
742+
const std::vector<int64_t>& shape = traced_input_shapes[i];
743+
const std::string& type = traced_input_types[i];
653744

654-
at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
655-
if (device == "gpu")
656-
t = t.cuda();
745+
at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
746+
if (device == "gpu")
747+
t = t.cuda();
657748

658-
input_tensors.push_back(t);
749+
input_tensors.push_back(t);
750+
}
659751
}
660752

661753
std::vector<at::Tensor> input_tensors2;
662-
for (size_t i = 0; i < input_shapes2.size(); i++)
754+
if (!input_contents2.empty())
663755
{
664-
const std::vector<int64_t>& shape = input_shapes2[i];
665-
const std::string& type = input_types2[i];
756+
if (input_contents2.size() != input_shapes2.size())
757+
{
758+
fprintf(stderr, "input2 expect %d tensors but got %d\n", (int)input_shapes2.size(), (int)input_contents2.size());
759+
return -1;
760+
}
666761

667-
at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
668-
if (device == "gpu")
669-
t = t.cuda();
762+
for (size_t i = 0; i < input_shapes2.size(); i++)
763+
{
764+
const std::vector<int64_t>& shape = input_shapes2[i];
765+
const std::string& type = input_types2[i];
766+
767+
size_t elem_count = 0;
768+
if (!shape_element_count(shape, elem_count))
769+
{
770+
fprintf(stderr, "invalid input2 shape for tensor %d\n", (int)i);
771+
return -1;
772+
}
670773

671-
input_tensors2.push_back(t);
774+
const size_t elemsize = input_type_to_elemsize(type);
775+
if (elemsize == 0)
776+
{
777+
fprintf(stderr, "unsupported input2 type %s\n", type.c_str());
778+
return -1;
779+
}
780+
781+
if (elem_count != input_contents2[i].size() / elemsize || elem_count * elemsize != input_contents2[i].size())
782+
{
783+
fprintf(stderr, "input2 tensor %d data size mismatch\n", (int)i);
784+
return -1;
785+
}
786+
787+
at::Tensor t = torch::from_blob((void*)input_contents2[i].data(), shape, torch::TensorOptions().dtype(input_type_to_c10_ScalarType(type))).clone();
788+
if (device == "gpu")
789+
t = t.cuda();
790+
791+
input_tensors2.push_back(t);
792+
}
793+
}
794+
else
795+
{
796+
for (size_t i = 0; i < input_shapes2.size(); i++)
797+
{
798+
const std::vector<int64_t>& shape = input_shapes2[i];
799+
const std::string& type = input_types2[i];
800+
801+
at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
802+
if (device == "gpu")
803+
t = t.cuda();
804+
805+
input_tensors2.push_back(t);
806+
}
672807
}
673808

674809
torch::jit::Module mod;

‎tools/pnnx/src/load_torchscript.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,10 @@ int load_torchscript(const std::string& ptpath, Graph& g,
1212
const std::string& device,
1313
const std::vector<std::vector<int64_t> >& input_shapes,
1414
const std::vector<std::string>& input_types,
15+
const std::vector<std::vector<char> >& input_contents,
1516
const std::vector<std::vector<int64_t> >& input_shapes2,
1617
const std::vector<std::string>& input_types2,
18+
const std::vector<std::vector<char> >& input_contents2,
1719
const std::vector<std::string>& customop_modules,
1820
const std::vector<std::string>& module_operators,
1921
const std::string& foldable_constants_zippath,

0 commit comments

Comments
 (0)