wehub-resource-sync
7254f7b4d1
Build / Build (macos-latest) (push) Has been cancelled
Build / Build (ubuntu-latest) (push) Has been cancelled
Build / Build (windows-latest) (push) Has been cancelled
Build / Analyze (javascript) (push) Has been cancelled
Build / Analyze (python) (push) Has been cancelled
593 lines
24 KiB
JavaScript
593 lines
24 KiB
JavaScript
|
|
import * as base from './base.js';
|
|
import * as flatbuffers from './flatbuffers.js';
|
|
import * as json from './json.js';
|
|
|
|
const rknn = {};
|
|
const openvx = {};
|
|
|
|
rknn.ModelFactory = class {
|
|
|
|
async match(context) {
|
|
const container = await rknn.Container.open(context);
|
|
if (container) {
|
|
return context.set('rknn', container);
|
|
}
|
|
return null;
|
|
}
|
|
|
|
async open(context) {
|
|
rknn.schema = await context.require('./rknn-schema');
|
|
rknn.schema = rknn.schema.rknn;
|
|
const metadata = await context.metadata('rknn-metadata.json');
|
|
const target = context.value;
|
|
target.read();
|
|
if (target.has('json')) {
|
|
const buffer = target.get('json');
|
|
const reader = json.TextReader.open(buffer);
|
|
const model = reader.read();
|
|
return new rknn.Model(metadata, 'json', model, target);
|
|
}
|
|
if (target.has('flatbuffers')) {
|
|
const buffer = target.get('flatbuffers');
|
|
const reader = flatbuffers.BinaryReader.open(buffer);
|
|
const model = rknn.schema.Model.create(reader);
|
|
return new rknn.Model(metadata, 'flatbuffers', model, null);
|
|
}
|
|
if (target.has('openvx')) {
|
|
const buffer = target.get('openvx');
|
|
const model = new openvx.Model(buffer);
|
|
return new rknn.Model(metadata, 'openvx', model, null);
|
|
}
|
|
throw new rknn.Error("Unsupported RKNN format.");
|
|
}
|
|
};
|
|
|
|
rknn.Model = class {
|
|
|
|
constructor(metadata, type, model, container) {
|
|
switch (type) {
|
|
case 'json': {
|
|
this.format = `RKNN v${model.version.split('-').shift()}`;
|
|
this.name = model.name || '';
|
|
this.producer = model.ori_network_platform || model.network_platform || '';
|
|
this.runtime = model.target_platform ? model.target_platform.join(',') : '';
|
|
this.modules = [new rknn.Graph(metadata, type, model.name || '', model, container)];
|
|
break;
|
|
}
|
|
case 'flatbuffers': {
|
|
const version = model.compiler.split('-').shift();
|
|
this.format = `RKNN Lite${version ? ` v${version}` : ''}`;
|
|
this.runtime = model.runtime;
|
|
this.name = model.name || '';
|
|
this.modules = model.graphs.map((graph) => new rknn.Graph(metadata, type, '', graph, null));
|
|
this.source = model.source;
|
|
break;
|
|
}
|
|
case 'openvx': {
|
|
this.format = 'RKNN OpenVX';
|
|
this.name = model.name || '';
|
|
this.modules = [new rknn.Graph(metadata, type, '', model, container)];
|
|
break;
|
|
}
|
|
default: {
|
|
throw new rknn.Error(`Unsupported RKNN model type '${type}'.`);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
rknn.Graph = class {
|
|
|
|
constructor(metadata, type, name, obj, container) {
|
|
this.name = name;
|
|
this.inputs = [];
|
|
this.outputs = [];
|
|
this.nodes = [];
|
|
switch (type) {
|
|
case 'json': {
|
|
const dataType = (value) => {
|
|
const type = value.vx_type.startsWith('VSI_NN_TYPE_') ? value.vx_type.split('_').pop().toLowerCase() : value.vx_type;
|
|
switch (type) {
|
|
case 'uint8':
|
|
case 'int8':
|
|
case 'int16':
|
|
case 'int32':
|
|
case 'int64':
|
|
case 'float16':
|
|
case 'float32':
|
|
case 'float64':
|
|
case 'vdata':
|
|
return type;
|
|
default:
|
|
if (value.vx_type !== '') {
|
|
throw new rknn.Error(`Invalid data type '${JSON.stringify(dataType)}'.`);
|
|
}
|
|
return '?';
|
|
}
|
|
};
|
|
const model = obj;
|
|
const values = new Map();
|
|
for (const const_tensor of model.const_tensor) {
|
|
const name = `const_tensor:${const_tensor.tensor_id}`;
|
|
const shape = new rknn.TensorShape(const_tensor.size);
|
|
if (const_tensor.data_type === 0) {
|
|
const value = new rknn.Value(name, null, null);
|
|
values.set(name, value);
|
|
} else {
|
|
const type = new rknn.TensorType(dataType(const_tensor.dtype), shape);
|
|
const tensor = new rknn.Tensor(type, const_tensor.offset, undefined, null);
|
|
const value = new rknn.Value(name, type, tensor);
|
|
values.set(name, value);
|
|
}
|
|
}
|
|
for (const virtual_tensor of model.virtual_tensor) {
|
|
const name = `${virtual_tensor.node_id}:${virtual_tensor.output_port}`;
|
|
const value = new rknn.Value(name, null, null);
|
|
values.set(name, value);
|
|
}
|
|
for (const norm_tensor of model.norm_tensor) {
|
|
const name = `norm_tensor:${norm_tensor.tensor_id}`;
|
|
const shape = new rknn.TensorShape(norm_tensor.size);
|
|
if (norm_tensor.dtype === 0) {
|
|
const value = new rknn.Value(name, null, null);
|
|
values.set(name, value);
|
|
} else {
|
|
const type = new rknn.TensorType(dataType(norm_tensor.dtype), shape);
|
|
const value = new rknn.Value(name, type, null);
|
|
values.set(name, value);
|
|
}
|
|
}
|
|
const value = (name) => {
|
|
if (!values.has(name)) {
|
|
values.set(name, new rknn.Value(name, null, null));
|
|
}
|
|
return values.get(name);
|
|
};
|
|
for (const node of model.nodes) {
|
|
node.input = [];
|
|
node.output = [];
|
|
}
|
|
for (const connection of model.connection) {
|
|
switch (connection.left) {
|
|
case 'input':
|
|
model.nodes[connection.node_id].input.push(connection);
|
|
if (connection.right_node) {
|
|
model.nodes[connection.right_node.node_id].output[connection.right_node.tensor_id] = connection;
|
|
}
|
|
break;
|
|
case 'output':
|
|
model.nodes[connection.node_id].output.push(connection);
|
|
break;
|
|
default:
|
|
throw new rknn.Error(`Unsupported left connection '${connection.left}'.`);
|
|
}
|
|
}
|
|
for (const graph of model.graph) {
|
|
const key = `${graph.right}:${graph.right_tensor_id}`;
|
|
const name = graph.left + (graph.left_tensor_id === 0 ? '' : graph.left_tensor_id.toString());
|
|
const argument = new rknn.Argument(name, [value(key)]);
|
|
switch (graph.left) {
|
|
case 'input':
|
|
this.inputs.push(argument);
|
|
break;
|
|
case 'output':
|
|
this.outputs.push(argument);
|
|
break;
|
|
default:
|
|
throw new rknn.Error(`Unsupported left graph connection '${graph.left}'.`);
|
|
}
|
|
}
|
|
this.nodes = model.nodes.map((node) => new rknn.Node(metadata, type, node, value, container));
|
|
break;
|
|
}
|
|
case 'flatbuffers': {
|
|
const graph = obj;
|
|
const dataTypes = ['?', 'float32', 'uint8', 'int8', 'uint16', 'int16', 'int32', 'int64', 'string', 'boolean', 'float16', 'float64', 'uint32', 'uint64', 'complex<float32>', 'complex<float64>', 'bfloat16'];
|
|
const args = graph.tensors.map((tensor) => {
|
|
const shape = new rknn.TensorShape(Array.from(tensor.shape));
|
|
const dataType = tensor.data_type < dataTypes.length ? dataTypes[tensor.data_type] : '?';
|
|
const type = new rknn.TensorType(dataType, shape);
|
|
const initializer = tensor.kind !== 4 && tensor.kind !== 5 ? null : new rknn.Tensor(type, 0, tensor.size, null);
|
|
return new rknn.Value(tensor.name, type, initializer);
|
|
});
|
|
const arg = (index) => {
|
|
if (index >= args.length) {
|
|
throw new rknn.Error(`Invalid tensor index '${index}'.`);
|
|
}
|
|
return args[index];
|
|
};
|
|
this.nodes = graph.nodes.map((node) => new rknn.Node(metadata, type, node, arg, container));
|
|
break;
|
|
}
|
|
case 'openvx': {
|
|
const model = obj;
|
|
this.nodes = model.nodes.map((node) => new rknn.Node(metadata, type, node, null, container));
|
|
break;
|
|
}
|
|
default: {
|
|
throw new rknn.Error(`Unsupported RKNN graph type '${type}'.`);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
rknn.Argument = class {
|
|
|
|
constructor(name, value) {
|
|
this.name = name;
|
|
this.value = value;
|
|
}
|
|
};
|
|
|
|
rknn.Value = class {
|
|
|
|
constructor(name, type = null, initializer = null) {
|
|
if (typeof name !== 'string') {
|
|
throw new rknn.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
|
|
}
|
|
this.name = name;
|
|
this.type = type;
|
|
this.initializer = initializer;
|
|
}
|
|
};
|
|
|
|
rknn.Node = class {
|
|
|
|
constructor(metadata, type, node, value, container) {
|
|
this.inputs = [];
|
|
this.outputs = [];
|
|
this.attributes = [];
|
|
switch (type) {
|
|
case 'json': {
|
|
this.name = node.name || '';
|
|
if (node.op === 'VSI_NN_OP_NBG' && container && container.has('openvx')) {
|
|
const buffer = container.get('openvx');
|
|
const model = new openvx.Model(buffer);
|
|
this.type = new rknn.Graph(metadata, 'openvx', 'NBG', model, null);
|
|
} else if (node.op === 'RKNN_OP_NNBG' && container && container.has('flatbuffers')) {
|
|
const buffer = container.get('flatbuffers');
|
|
const reader = flatbuffers.BinaryReader.open(buffer);
|
|
const model = rknn.schema.Model.create(reader);
|
|
this.type = new rknn.Graph(metadata, 'flatbuffers', 'NNBG', model.graphs[0], null);
|
|
} else {
|
|
const type = metadata.type(node.op);
|
|
this.type = type ? { ...type } : { name: node.op };
|
|
for (const prefix of ['VSI_NN_OP_', 'RKNN_OP_']) {
|
|
this.type.name = this.type.name.startsWith(prefix) ? this.type.name.substring(prefix.length) : this.type.name;
|
|
}
|
|
}
|
|
node.input = node.input || [];
|
|
for (let i = 0; i < node.input.length;) {
|
|
const input = this.type && this.type.inputs && i < this.type.inputs.length ? this.type.inputs[i] : { name: i === 0 ? 'input' : i.toString() };
|
|
const count = input.list ? node.input.length - i : 1;
|
|
const list = node.input.slice(i, i + count).map((input) => {
|
|
if (input.right_tensor) {
|
|
return value(`${input.right_tensor.type}:${input.right_tensor.tensor_id}`);
|
|
}
|
|
if (input.right_node) {
|
|
return value(`${input.right_node.node_id}:${input.right_node.tensor_id}`);
|
|
}
|
|
throw new rknn.Error('Invalid input argument.');
|
|
});
|
|
this.inputs.push(new rknn.Argument(input.name, list));
|
|
i += count;
|
|
}
|
|
node.output = node.output || [];
|
|
for (let i = 0; i < node.output.length;) {
|
|
const output = this.type.outputs && i < this.type.outputs.length ? this.type.outputs[i] : { name: i === 0 ? 'output' : i.toString() };
|
|
const count = output.list ? node.output.length - i : 1;
|
|
const list = node.output.slice(i, i + count).map((output) => {
|
|
if (output.right_tensor) {
|
|
return value(`${output.right_tensor.type}:${output.right_tensor.tensor_id}`);
|
|
}
|
|
if (output.right_node) {
|
|
return value(`${output.right_node.node_id}:${output.right_node.tensor_id}`);
|
|
}
|
|
throw new rknn.Error('Invalid output argument.');
|
|
});
|
|
this.outputs.push(new rknn.Argument(output.name, list));
|
|
i += count;
|
|
}
|
|
if (node.nn) {
|
|
for (const params of Object.values(node.nn)) {
|
|
for (const [name, value] of Object.entries(params)) {
|
|
const attribute = new rknn.Argument(name, value);
|
|
this.attributes.push(attribute);
|
|
}
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case 'flatbuffers': {
|
|
this.name = node.name;
|
|
this.type = metadata.type(node.type);
|
|
if (node.inputs.length > 0) {
|
|
const inputs = this.type.inputs || (node.inputs.length === 1 ? [{ name: "input" }] : [{ name: "inputs", list: true }]);
|
|
if (Array.isArray(inputs) && inputs.length > 0 && inputs[0].list === true) {
|
|
this.inputs = [new rknn.Argument(inputs[0].name, Array.from(node.inputs).map((input) => value(input)))];
|
|
} else {
|
|
this.inputs = Array.from(node.inputs).map((input, index) => {
|
|
return new rknn.Argument(index < inputs.length ? inputs[index].name : index.toString(), [value(input)]);
|
|
});
|
|
}
|
|
}
|
|
if (node.outputs.length > 0) {
|
|
const outputs = this.type.outputs || (node.outputs.length === 1 ? [{ name: "output" }] : [{ name: "outputs", list: true }]);
|
|
if (Array.isArray(outputs) && outputs.length > 0 && outputs[0].list === true) {
|
|
const values = Array.from(node.outputs).map((output) => value(output));
|
|
const argument = new rknn.Argument(outputs[0].name, values);
|
|
this.outputs = [argument];
|
|
} else {
|
|
this.outputs = Array.from(node.outputs).map((output, index) => {
|
|
return new rknn.Argument(index < outputs.length ? outputs[index].name : index.toString(), [value(output)]);
|
|
});
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case 'openvx': {
|
|
this.name = '';
|
|
this.type = metadata.type(node.type);
|
|
break;
|
|
}
|
|
default: {
|
|
throw new rknn.Error(`Unsupported RKNN node type '${type}'.`);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
rknn.Tensor = class {
|
|
|
|
constructor(type, offset, size, weights) {
|
|
this.type = type;
|
|
this.values = null;
|
|
let itemsize = 0;
|
|
switch (this.type.dataType) {
|
|
case 'uint8': itemsize = 1; break;
|
|
case 'int8': itemsize = 1; break;
|
|
case 'int16': itemsize = 2; break;
|
|
case 'int32': itemsize = 4; break;
|
|
case 'int64': itemsize = 8; break;
|
|
case 'uint16': itemsize = 2; break;
|
|
case 'uint32': itemsize = 4; break;
|
|
case 'uint64': itemsize = 8; break;
|
|
case 'float16': itemsize = 2; break;
|
|
case 'bfloat16': itemsize = 2; break;
|
|
case 'float32': itemsize = 4; break;
|
|
case 'float64': itemsize = 8; break;
|
|
case 'boolean': itemsize = 1; break;
|
|
case 'vdata': itemsize = 1; break;
|
|
case 'string': itemsize = 1; break;
|
|
case '?': itemsize = 0; break;
|
|
default: throw new rknn.Error(`Unsupported tensor data type '${this.type.dataType}'.`);
|
|
}
|
|
if (weights) {
|
|
const shape = type.shape.dimensions;
|
|
const count = shape.reduce((a, b) => a * b, 1);
|
|
const length = itemsize * count;
|
|
if (length > 0) {
|
|
if (size !== undefined && size !== length) {
|
|
throw new rknn.Error(`Tensor size mismatch for '${this.type.dataType}'. Expected '${length}' bytes but got '${size}' bytes.`);
|
|
}
|
|
this.values = weights.slice(offset, offset + length);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
rknn.TensorType = class {
|
|
|
|
constructor(dataType, shape) {
|
|
this.dataType = dataType;
|
|
this.shape = shape;
|
|
}
|
|
|
|
toString() {
|
|
return this.dataType + this.shape.toString();
|
|
}
|
|
};
|
|
|
|
rknn.TensorShape = class {
|
|
|
|
constructor(dimensions) {
|
|
this.dimensions = dimensions;
|
|
}
|
|
|
|
toString() {
|
|
if (!this.dimensions || this.dimensions.length === 0) {
|
|
return '';
|
|
}
|
|
return `[${this.dimensions.join(',')}]`;
|
|
}
|
|
};
|
|
|
|
rknn.Container = class extends Map {
|
|
|
|
static async open(context) {
|
|
const stream = context.stream;
|
|
if (stream) {
|
|
const signature = rknn.Container.signature(stream);
|
|
switch (signature) {
|
|
case 'rknn':
|
|
case 'openvx':
|
|
case 'flatbuffers':
|
|
case 'cyptrknn':
|
|
return new rknn.Container(stream, signature);
|
|
default:
|
|
break;
|
|
}
|
|
const obj = await context.peek('json');
|
|
if (obj && obj.version && Array.isArray(obj.nodes) && obj.network_platform) {
|
|
const entries = new Map();
|
|
entries.set('json', stream);
|
|
return new rknn.Container(null, null, entries);
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
constructor(stream, signature, entries) {
|
|
super(entries);
|
|
this.stream = stream;
|
|
this.signature = signature;
|
|
}
|
|
|
|
read() {
|
|
const stream = this.stream;
|
|
if (stream) {
|
|
switch (this.signature) {
|
|
case 'rknn': {
|
|
const uint64 = () => {
|
|
const buffer = stream.read(8);
|
|
const reader = base.BinaryReader.open(buffer);
|
|
return reader.uint64();
|
|
};
|
|
stream.skip(8);
|
|
const version = uint64();
|
|
if ((version >> 8n) !== 0n && (version >> 8n) !== 0x10n) {
|
|
throw new rknn.Error(`Unsupported RKNN container version '${version}'.`);
|
|
}
|
|
const data_size = uint64().toNumber();
|
|
if ((version & 0xffn) > 1n && data_size > 0) {
|
|
stream.skip(40);
|
|
}
|
|
const signature = rknn.Container.signature(stream, data_size);
|
|
const data = stream.read(data_size);
|
|
const json_size = uint64().toNumber();
|
|
const json = stream.read(json_size);
|
|
this.set('json', json);
|
|
if (signature) {
|
|
this.set(signature, data);
|
|
}
|
|
break;
|
|
}
|
|
case 'openvx':
|
|
case 'flatbuffers': {
|
|
this.set(this.signature, stream.peek());
|
|
break;
|
|
}
|
|
case 'cyptrknn': {
|
|
throw new rknn.Error('Invalid file content. File contains undocumented encrypted RKNN data.');
|
|
}
|
|
default: {
|
|
break;
|
|
}
|
|
}
|
|
delete this.stream;
|
|
}
|
|
}
|
|
|
|
static signature(stream, length) {
|
|
length = length || stream.length;
|
|
if (stream && (stream.position + 16) <= length) {
|
|
const signature = [0x52, 0x4B, 0x4E, 0x4E]; // RKNN
|
|
if (stream.peek(signature.length).every((value, index) => value === signature[index])) {
|
|
return 'rknn';
|
|
}
|
|
}
|
|
if (stream && (stream.position + 16) <= length) {
|
|
const signature = [0x43, 0x59, 0x50, 0x54, 0x52, 0x4B, 0x4E, 0x4E]; // CYPTRKNN
|
|
if (stream.peek(signature.length).every((value, index) => value === signature[index])) {
|
|
return 'cyptrknn';
|
|
}
|
|
}
|
|
if (stream && (stream.position + 8) <= length) {
|
|
const signature = [0x52, 0x4B, 0x4E, 0x4E]; // RKNN
|
|
if (stream.peek(8).subarray(4, 8).every((value, index) => value === signature[index])) {
|
|
return 'flatbuffers';
|
|
}
|
|
}
|
|
if (stream && (stream.position + 8) <= length) {
|
|
const signature = [0x56, 0x50, 0x4D, 0x4E]; // VPMN
|
|
if (stream.peek(signature.length).every((value, index) => value === signature[index])) {
|
|
return 'openvx';
|
|
}
|
|
}
|
|
return undefined;
|
|
}
|
|
};
|
|
|
|
openvx.BufferReader = class {
|
|
|
|
constructor(buffer) {
|
|
this._reader = base.BinaryReader.open(buffer);
|
|
}
|
|
|
|
seek(position) {
|
|
this._reader.seek(position);
|
|
}
|
|
|
|
skip(offset) {
|
|
this._reader.skip(offset);
|
|
}
|
|
|
|
read(length) {
|
|
return this._reader.read(length);
|
|
}
|
|
|
|
uint16() {
|
|
return this._reader.uint16();
|
|
}
|
|
|
|
uint32() {
|
|
return this._reader.uint32();
|
|
}
|
|
|
|
string(length) {
|
|
const buffer = this.read(length);
|
|
const index = buffer.indexOf(0);
|
|
const data = index === -1 ? buffer : buffer.subarray(0, index);
|
|
this._decoder = this._decoder || new TextDecoder('ascii');
|
|
return this._decoder.decode(data);
|
|
}
|
|
};
|
|
|
|
openvx.Model = class {
|
|
|
|
constructor(buffer) {
|
|
const reader = new openvx.BufferReader(buffer);
|
|
reader.skip(4); // signature
|
|
const major = reader.uint16();
|
|
/* const minor = */ reader.uint16();
|
|
reader.skip(4);
|
|
this.name = reader.string(64);
|
|
this.nodes = new Array(reader.uint32());
|
|
if (major > 3) {
|
|
reader.skip(296);
|
|
} else if (major > 1) {
|
|
reader.skip(288);
|
|
} else {
|
|
reader.skip(32);
|
|
}
|
|
/* const inputOffset = */ reader.uint32();
|
|
/* const inputSize = */ reader.uint32();
|
|
/* const outputOffset = */ reader.uint32();
|
|
/* const outputSize = */ reader.uint32();
|
|
const nodeOffset = reader.uint32();
|
|
/* const nodeSize = */ reader.uint32();
|
|
reader.seek(nodeOffset);
|
|
for (let i = 0; i < this.nodes.length; i++) {
|
|
const type = reader.string(64);
|
|
const node = { type };
|
|
node.index = reader.uint32();
|
|
node.c = reader.uint32();
|
|
if (major > 3) {
|
|
node.d = reader.uint32();
|
|
}
|
|
this.nodes[i] = node;
|
|
}
|
|
}
|
|
};
|
|
|
|
rknn.Error = class extends Error {
|
|
|
|
constructor(message) {
|
|
super(message);
|
|
this.name = 'Error loading RKNN model.';
|
|
}
|
|
};
|
|
|
|
export const ModelFactory = rknn.ModelFactory;
|