Skip to content

pnnx: support npy input tensors #6700

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nihui merged 2 commits into from
May 6, 2026
Merged

pnnx: support npy input tensors #6700

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dc600e0
pnnx: support npy input tensors
nihui May 6, 2026
182b148
Merge branch 'master' into pnnx-npy
nihui May 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions tools/pnnx/src/ir.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2771,14 +2771,14 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
fprintf(pyfp, " torch.manual_seed(0)\n");

int input_shapes_i = 0;

std::vector<std::string> input_names;
for (const Operator* op : ops)
{
if (op->type != "pnnx.Input")
continue;

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

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

std::string input_name = std::string("v_") + sanitize_identifier(r->name);
if (type_is_integer(r->type))
{
fprintf(pyfp, " %s = torch.randint(10, (", input_name.c_str());
Expand Down
65 changes: 63 additions & 2 deletions tools/pnnx/src/load_tnn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,26 @@ static int get_tnn_tensor_type(int dt)
return 0; // unknown type
}

static int input_type_to_pnnx_type(const std::string& t)
{
if (t == "f32") return 1;
if (t == "f64") return 2;
if (t == "f16") return 3;
if (t == "i32") return 4;
if (t == "i64") return 5;
if (t == "i16") return 6;
if (t == "i8") return 7;
if (t == "u8") return 8;
if (t == "bool") return 9;
if (t == "c64") return 10;
if (t == "c128") return 11;
if (t == "c32") return 12;
if (t == "bf16") return 13;

fprintf(stderr, "unsupported input type %s\n", t.c_str());
return 0;
}

Attribute::Attribute(FILE* bp)
{
unsigned int magic;
Expand Down Expand Up @@ -178,7 +198,9 @@ Attribute::Attribute(FILE* bp)
fread((void*)data.data(), 1, length, bp);
}

int load_tnn(const std::string& tnnpath, Graph& pnnx_graph)
int load_tnn(const std::string& tnnpath, Graph& pnnx_graph,
const std::vector<std::vector<int64_t> >& input_shapes,
const std::vector<std::string>& input_types)
{
fprintf(stderr, "############# pass_level0 tnn\n");

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

int tensor_type = get_tnn_tensor_type(datatype);

if (!input_shapes.empty())
{
if (input_shapes.size() != 1)
{
fprintf(stderr, "tnn input expect 1 tensor but got %d\n", (int)input_shapes.size());
return -1;
}
if (input_types.size() != input_shapes.size())
{
fprintf(stderr, "tnn input type count mismatch, expect %d but got %d\n", (int)input_shapes.size(), (int)input_types.size());
return -1;
}

shape.clear();
const std::vector<int64_t>& input_shape = input_shapes[0];
for (size_t i = 0; i < input_shape.size(); i++)
{
if (input_shape[i] < 0 || input_shape[i] > INT_MAX)
{
fprintf(stderr, "invalid tnn input shape dimension %ld\n", input_shape[i]);
return -1;
}
shape.push_back((int)input_shape[i]);
}

const int input_tensor_type = input_type_to_pnnx_type(input_types[0]);
if (input_tensor_type == 0)
return -1;
if (tensor_type != 0 && tensor_type != input_tensor_type)
{
Comment thread
nihui marked this conversation as resolved.
fprintf(stderr, "tnn input type mismatch, model type %d but got %s\n", tensor_type, input_types[0].c_str());
return -1;
}

tensor_type = input_tensor_type;
}

Operator* op = pnnx_graph.new_operator("pnnx.Input", "input0");

Operand* r = pnnx_graph.new_operand(blob_name);

r->producer = op;

r->shape = shape;
r->type = get_tnn_tensor_type(datatype);
r->type = tensor_type;

op->outputs.push_back(r);
}
Expand Down
4 changes: 3 additions & 1 deletion tools/pnnx/src/load_tnn.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@

namespace pnnx {

int load_tnn(const std::string& tnnpath, Graph& g);
int load_tnn(const std::string& tnnpath, Graph& g,
const std::vector<std::vector<int64_t> >& input_shapes,
const std::vector<std::string>& input_types);

} // namespace pnnx

Expand Down
163 changes: 149 additions & 14 deletions tools/pnnx/src/load_torchscript.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
#include <torch/script.h>
#include <torch/csrc/api/include/torch/version.h>
#include <torch/csrc/jit/serialization/import_read.h>

#include <limits>
#ifdef PNNX_TORCHVISION
namespace vision {
int64_t cuda_version();
Expand Down Expand Up @@ -418,11 +420,48 @@ static c10::ScalarType input_type_to_c10_ScalarType(const std::string& t)
if (t == "i64") return torch::kInt64;
if (t == "i8") return torch::kInt8;
if (t == "u8") return torch::kUInt8;
if (t == "bool") return torch::kBool;

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

static size_t input_type_to_elemsize(const std::string& t)
{
if (t == "f32") return 4;
if (t == "f64") return 8;
if (t == "f16") return 2;
if (t == "i32") return 4;
if (t == "i64") return 8;
if (t == "i16") return 2;
if (t == "i8") return 1;
if (t == "u8") return 1;
if (t == "bool") return 1;
if (t == "c64") return 8;
if (t == "c128") return 16;
if (t == "c32") return 4;
if (t == "bf16") return 2;
return 0;
}

static bool shape_element_count(const std::vector<int64_t>& shape, size_t& elem_count)
{
elem_count = 1;
for (size_t i = 0; i < shape.size(); i++)
{
if (shape[i] < 0)
return false;

const size_t dim = (size_t)shape[i];
if (dim != 0 && elem_count > std::numeric_limits<size_t>::max() / dim)
return false;

elem_count *= dim;
}

return true;
}

static const char* get_at_tensor_type_str(const at::ScalarType& st)
{
if (st == c10::ScalarType::Float) return "f32";
Expand All @@ -437,6 +476,7 @@ static const char* get_at_tensor_type_str(const at::ScalarType& st)
if (st == c10::ScalarType::ComplexDouble) return "c128";
if (st == c10::ScalarType::ComplexHalf) return "c32";
if (st == c10::ScalarType::BFloat16) return "bf16";
if (st == c10::ScalarType::Bool) return "bool";

// unknown
fprintf(stderr, "unsupported tensor elem data type %d\n", (int)st);
Expand Down Expand Up @@ -583,8 +623,10 @@ int load_torchscript(const std::string& ptpath, Graph& pnnx_graph,
const std::string& device,
const std::vector<std::vector<int64_t> >& input_shapes,
const std::vector<std::string>& input_types,
const std::vector<std::vector<char> >& input_contents,
const std::vector<std::vector<int64_t> >& input_shapes2,
const std::vector<std::string>& input_types2,
const std::vector<std::vector<char> >& input_contents2,
const std::vector<std::string>& customop_modules,
const std::vector<std::string>& module_operators,
const std::string& foldable_constants_zippath,
Expand Down Expand Up @@ -646,29 +688,122 @@ int load_torchscript(const std::string& ptpath, Graph& pnnx_graph,
}

std::vector<at::Tensor> input_tensors;
for (size_t i = 0; i < traced_input_shapes.size(); i++)
if (!input_contents.empty())
{
if (input_contents.size() != traced_input_shapes.size())
{
fprintf(stderr, "input expect %d tensors but got %d\n", (int)traced_input_shapes.size(), (int)input_contents.size());
return -1;
}

for (size_t i = 0; i < traced_input_shapes.size(); i++)
{
const std::vector<int64_t>& shape = traced_input_shapes[i];
const std::string& type = traced_input_types[i];

size_t elem_count = 0;
if (!shape_element_count(shape, elem_count))
{
fprintf(stderr, "invalid input shape for tensor %d\n", (int)i);
return -1;
}

const size_t elemsize = input_type_to_elemsize(type);
if (elemsize == 0)
{
fprintf(stderr, "unsupported input type %s\n", type.c_str());
return -1;
}

if (elem_count != input_contents[i].size() / elemsize || elem_count * elemsize != input_contents[i].size())
{
fprintf(stderr, "input tensor %d data size mismatch for shape [", (int)i);
for (size_t j = 0; j < shape.size(); j++)
{
fprintf(stderr, "%ld", shape[j]);
if (j + 1 != shape.size())
fprintf(stderr, ",");
}
fprintf(stderr, "]%s\n", type.c_str());
return -1;
}

at::Tensor t = torch::from_blob((void*)input_contents[i].data(), shape, torch::TensorOptions().dtype(input_type_to_c10_ScalarType(type))).clone();
if (device == "gpu")
t = t.cuda();

input_tensors.push_back(t);
}
}
else
{
const std::vector<int64_t>& shape = traced_input_shapes[i];
const std::string& type = traced_input_types[i];
for (size_t i = 0; i < traced_input_shapes.size(); i++)
{
const std::vector<int64_t>& shape = traced_input_shapes[i];
const std::string& type = traced_input_types[i];

at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
if (device == "gpu")
t = t.cuda();
at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
if (device == "gpu")
t = t.cuda();

input_tensors.push_back(t);
input_tensors.push_back(t);
}
}

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

at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
if (device == "gpu")
t = t.cuda();
for (size_t i = 0; i < input_shapes2.size(); i++)
{
const std::vector<int64_t>& shape = input_shapes2[i];
const std::string& type = input_types2[i];

size_t elem_count = 0;
if (!shape_element_count(shape, elem_count))
{
fprintf(stderr, "invalid input2 shape for tensor %d\n", (int)i);
return -1;
}

input_tensors2.push_back(t);
const size_t elemsize = input_type_to_elemsize(type);
if (elemsize == 0)
{
fprintf(stderr, "unsupported input2 type %s\n", type.c_str());
return -1;
}

if (elem_count != input_contents2[i].size() / elemsize || elem_count * elemsize != input_contents2[i].size())
{
fprintf(stderr, "input2 tensor %d data size mismatch\n", (int)i);
return -1;
}

at::Tensor t = torch::from_blob((void*)input_contents2[i].data(), shape, torch::TensorOptions().dtype(input_type_to_c10_ScalarType(type))).clone();
if (device == "gpu")
t = t.cuda();

input_tensors2.push_back(t);
}
}
else
{
for (size_t i = 0; i < input_shapes2.size(); i++)
{
const std::vector<int64_t>& shape = input_shapes2[i];
const std::string& type = input_types2[i];

at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
if (device == "gpu")
t = t.cuda();

input_tensors2.push_back(t);
}
}

torch::jit::Module mod;
Expand Down
2 changes: 2 additions & 0 deletions tools/pnnx/src/load_torchscript.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,10 @@ int load_torchscript(const std::string& ptpath, Graph& g,
const std::string& device,
const std::vector<std::vector<int64_t> >& input_shapes,
const std::vector<std::string>& input_types,
const std::vector<std::vector<char> >& input_contents,
const std::vector<std::vector<int64_t> >& input_shapes2,
const std::vector<std::string>& input_types2,
const std::vector<std::vector<char> >& input_contents2,
const std::vector<std::string>& customop_modules,
const std::vector<std::string>& module_operators,
const std::string& foldable_constants_zippath,
Expand Down
Loading
Loading