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
634 lines
25 KiB
JavaScript
634 lines
25 KiB
JavaScript
|
|
const tnn = {};
|
|
|
|
tnn.ModelFactory = class {
|
|
|
|
async match(context) {
|
|
const identifier = context.identifier.toLowerCase();
|
|
const stream = context.stream;
|
|
if (stream && identifier.endsWith('.tnnproto')) {
|
|
try {
|
|
const reader = await context.read('text', 0x10000);
|
|
const content = reader.read('\n');
|
|
if (content !== undefined) {
|
|
const line = content.trim();
|
|
if (line.startsWith('"') && line.endsWith('"')) {
|
|
const header = line.replace(/(^")|("$)/g, '').split(',').shift().trim().split(' ');
|
|
if (header.length === 3 || (header.length >= 4 && (header[3] === '4206624770' || header[3] === '4206624772'))) {
|
|
return context.set('tnn.model');
|
|
}
|
|
}
|
|
}
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
}
|
|
if (stream && identifier.endsWith('.tnnmodel')) {
|
|
for (const signature of [[0x02, 0x00, 0xbc, 0xfa], [0x04, 0x00, 0xbc, 0xfa]]) {
|
|
if (signature.length <= stream.length && stream.peek(signature.length).every((value, index) => value === signature[index])) {
|
|
return context.set('tnn.params');
|
|
}
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
async open(context) {
|
|
const metadata = await context.metadata('tnn-metadata.json');
|
|
switch (context.type) {
|
|
case 'tnn.model': {
|
|
const name = `${context.identifier.substring(0, context.identifier.length - 9)}.tnnmodel`;
|
|
const reader = await context.read('text');
|
|
try {
|
|
const content = await context.fetch(name);
|
|
const resources = await tnn.LayerResourceReader.open(content);
|
|
return new tnn.Model(metadata, reader, resources);
|
|
} catch {
|
|
const resources = await tnn.LayerResourceReader.open(null);
|
|
return new tnn.Model(metadata, reader, resources);
|
|
}
|
|
}
|
|
case 'tnn.params': {
|
|
const name = `${context.identifier.substring(0, context.identifier.length - 9)}.tnnproto`;
|
|
const content = await context.fetch(name, null);
|
|
const reader = await content.read('text');
|
|
const resources = await tnn.LayerResourceReader.open(context);
|
|
return new tnn.Model(metadata, reader, resources);
|
|
}
|
|
default: {
|
|
throw new tnn.Error(`Unsupported TNN format '${context.type}'.`);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
tnn.Model = class {
|
|
|
|
constructor(metadata, tnnproto, resources) {
|
|
this.format = 'TNN';
|
|
this.modules = [new tnn.Graph(metadata, tnnproto, resources)];
|
|
}
|
|
};
|
|
|
|
tnn.Graph = class {
|
|
|
|
constructor(metadata, tnnproto, resources) {
|
|
this.inputs = [];
|
|
this.outputs = [];
|
|
this.nodes = [];
|
|
const reader = new tnn.TextProtoReader(tnnproto);
|
|
reader.read('\n');
|
|
const values = new Map();
|
|
values.map = (name, type, tensor) => {
|
|
if (name.length === 0) {
|
|
return new tnn.Value(name, type || null, tensor || null);
|
|
}
|
|
if (!values.has(name)) {
|
|
values.set(name, new tnn.Value(name, type || null, tensor || null));
|
|
} else if (type || tensor) {
|
|
throw new tnn.Value(`Duplicate value '${name}'.`);
|
|
}
|
|
return values.get(name);
|
|
};
|
|
for (const input of reader.inputs) {
|
|
const shape = new tnn.TensorShape(input.shape);
|
|
const type = new tnn.TensorType(input.data_type, shape);
|
|
const argument = new tnn.Argument(input.name, [values.map(input.name, type)]);
|
|
this.inputs.push(argument);
|
|
}
|
|
for (const output of reader.outputs) {
|
|
const argument = new tnn.Argument(output.name, [values.map(output.name)]);
|
|
this.outputs.push(argument);
|
|
}
|
|
for (const layer of reader.layers) {
|
|
const node = new tnn.Node(metadata, resources, layer, values);
|
|
this.nodes.push(node);
|
|
}
|
|
}
|
|
};
|
|
|
|
tnn.Argument = class {
|
|
|
|
constructor(name, value, type = null, visible = true) {
|
|
this.name = name;
|
|
this.value = value;
|
|
this.type = type;
|
|
this.visible = visible;
|
|
}
|
|
};
|
|
|
|
tnn.Value = class {
|
|
|
|
constructor(name, type, initializer = null) {
|
|
if (typeof name !== 'string') {
|
|
throw new tnn.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
|
|
}
|
|
this.name = name;
|
|
this.type = initializer ? initializer.type : type;
|
|
this.initializer = initializer;
|
|
}
|
|
};
|
|
|
|
tnn.Node = class {
|
|
|
|
constructor(metadata, resources, layer, values) {
|
|
this.inputs = [];
|
|
this.outputs = [];
|
|
this.attributes = [];
|
|
this.name = layer.name;
|
|
this.type = { ...metadata.type(layer.type) };
|
|
delete this.type.identifier;
|
|
const entries = Array.from(layer.params);
|
|
for (let i = 0; i < entries.length;) {
|
|
const metadata = this.type && Array.isArray(this.type.attributes) ? this.type.attributes[i] : null;
|
|
let name = '';
|
|
let value = null;
|
|
let type = '';
|
|
let visible = true;
|
|
if (metadata && metadata.type === 'int32[]' && metadata.size) {
|
|
const size = parseInt(layer.params.get(metadata.size), 10);
|
|
value = entries.slice(i, i + size).map(([, value]) => parseInt(value, 10));
|
|
i += size;
|
|
} else {
|
|
[name, value] = entries[i];
|
|
i += 1;
|
|
}
|
|
if (metadata) {
|
|
name = metadata.name ? metadata.name : name;
|
|
type = metadata.type ? metadata.type : type;
|
|
switch (type) {
|
|
case '':
|
|
break;
|
|
case 'int32':
|
|
value = parseInt(value, 10);
|
|
break;
|
|
case 'float32':
|
|
value = parseFloat(value);
|
|
break;
|
|
case 'int32[]':
|
|
value = value.map((v) => parseInt(v, 10));
|
|
break;
|
|
default:
|
|
throw new tnn.Error(`Unsupported attribute type '${type}'.`);
|
|
}
|
|
visible = (metadata.visible === false) || (metadata.default !== undefined && (value === metadata.default || (value && value.toString() === metadata.default.toString()))) ? false : visible;
|
|
}
|
|
const argument = new tnn.Argument(name, value, type, visible);
|
|
this.attributes.push(argument);
|
|
}
|
|
const inputs = layer.inputs;
|
|
let inputIndex = 0;
|
|
if (this.type && this.type.inputs) {
|
|
for (const inputDef of this.type.inputs) {
|
|
if (inputIndex < inputs.length || inputDef.option !== 'optional') {
|
|
const inputCount = (inputDef.type === 'Tensor[]') ? (inputs.length - inputIndex) : 1;
|
|
const inputArguments = inputs.slice(inputIndex, inputIndex + inputCount).filter((id) => id !== '' || inputDef.option !== 'optional').map((id) => values.map(id));
|
|
const argument = new tnn.Argument(inputDef.name, inputArguments);
|
|
this.inputs.push(argument);
|
|
inputIndex += inputCount;
|
|
}
|
|
}
|
|
} else {
|
|
this.inputs.push(...inputs.slice(inputIndex).map((input, index) => {
|
|
const inputName = ((inputIndex + index) === 0) ? 'input' : (inputIndex + index).toString();
|
|
return new tnn.Argument(inputName, [values.map(input)]);
|
|
}));
|
|
}
|
|
const outputs = layer.outputs;
|
|
let outputIndex = 0;
|
|
if (this.type && this.type.outputs) {
|
|
for (const outputDef of this.type.outputs) {
|
|
if (outputIndex < outputs.length || outputDef.option !== 'optional') {
|
|
const outputCount = (outputDef.option === 'variadic') ? (outputs.length - outputIndex) : 1;
|
|
const outputArguments = outputs.slice(outputIndex, outputIndex + outputCount).map((id) => values.map(id));
|
|
const argument = new tnn.Argument(outputDef.name, outputArguments);
|
|
this.outputs.push(argument);
|
|
outputIndex += outputCount;
|
|
}
|
|
}
|
|
} else {
|
|
this.outputs.push(...outputs.slice(outputIndex).map((output, index) => {
|
|
const outputName = ((outputIndex + index) === 0) ? 'output' : (outputIndex + index).toString();
|
|
return new tnn.Argument(outputName, [values.map(output)]);
|
|
}));
|
|
}
|
|
const weight = (resource, name, shape) => {
|
|
const initializer = resource[name];
|
|
if (!initializer) {
|
|
throw new tnn.Error(`Layer initializer'${resource.type}.${name}' not found '`);
|
|
}
|
|
const tensor = new tnn.Tensor(new tnn.TensorType(initializer.dataType, new tnn.TensorShape(shape)), initializer.value);
|
|
const argument = new tnn.Argument(name, [values.map('', null, tensor)]);
|
|
this.inputs.push(argument);
|
|
};
|
|
const params = layer.params;
|
|
switch (this.type.name) {
|
|
case 'Convolution':
|
|
case 'ConvolutionDepthWise':
|
|
case 'Deconvolution':
|
|
case 'DeconvolutionDepthWise': {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
const num_output = parseInt(params.get('2') || 0, 10);
|
|
const kernel_w = parseInt(params.get('3') || 0, 10);
|
|
const kernel_h = parseInt(params.get('4') || kernel_w, 10);
|
|
const weight_data_size = resource.filter.length;
|
|
weight(resource, 'filter', [num_output, weight_data_size / (num_output * kernel_w * kernel_h), kernel_w, kernel_h]);
|
|
if (resource.bias) {
|
|
weight(resource, 'bias', [num_output]);
|
|
}
|
|
if (resource.quantized) {
|
|
weight(resource, 'quantized', [num_output]);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case 'Conv3D':{
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
const num_output = parseInt(params.get('2') || 0, 10);
|
|
const kernel_w = parseInt(params.get('3') || 0, 10);
|
|
const kernel_h = parseInt(params.get('4') || kernel_w, 10);
|
|
const kernel_d = parseInt(params.get('5') || kernel_w, 10);
|
|
const weight_data_size = resource.filter.length;
|
|
weight(resource, 'weight', [num_output, weight_data_size / (num_output * kernel_w * kernel_h * kernel_d), kernel_w, kernel_h, kernel_d]);
|
|
if (resource.bias) {
|
|
weight(resources, 'bias', [num_output]);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case 'InnerProduct': {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
const num_output = parseInt(params.get('0') || 0, 10);
|
|
const weight_data_size = resource.weight.length;
|
|
weight(resource, 'weight', [num_output, weight_data_size / num_output]);
|
|
weight(resource, 'bias', [num_output]);
|
|
if (resource.weight.dataType === 'int8') {
|
|
weight(resource, 'scale', [num_output]);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case 'PReLU': {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
weight(resource, 'slope', [resource.slope.length]);
|
|
}
|
|
break;
|
|
}
|
|
case 'BatchNormCxx':
|
|
case 'InstBatchNormCxx': {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
weight(resource, 'scale', [resource.scale.length]);
|
|
weight(resource, 'bias', [resource.bias.length]);
|
|
}
|
|
break;
|
|
}
|
|
case 'Div':
|
|
case 'Sub':
|
|
case 'Add':
|
|
case 'Mul':
|
|
case 'MatMul': {
|
|
if (this.inputs.length === 1) {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
const num_output = resource.slope.length;
|
|
weight(resource, 'slope', [num_output]);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case 'HdrGuide': {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
const weight_size = resource.ccm_weight.length;
|
|
weight(resource, 'ccm_weight', [weight_size]);
|
|
weight(resource, 'ccm_bias', [weight_size]);
|
|
weight(resource, 'shifts', [weight_size]);
|
|
weight(resource, 'slopes', [weight_size]);
|
|
weight(resource, 'projection_weight', [weight_size]);
|
|
weight(resource, 'projection_bias', [weight_size]);
|
|
}
|
|
break;
|
|
}
|
|
case 'BlobScale': {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
const scale_data_size = resource.scale.length;
|
|
weight(resource, 'scale', [scale_data_size]);
|
|
weight(resource, 'bias', [scale_data_size]);
|
|
}
|
|
break;
|
|
}
|
|
case 'Gather': {
|
|
const resource = resources.get(this.name);
|
|
if (resource) {
|
|
if (resource.data) {
|
|
weight(resource, 'data', [resource.data.length]);
|
|
}
|
|
if (resource.indices) {
|
|
weight(resource, 'indices', [resource.indices.length]);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
default: {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
tnn.Tensor = class {
|
|
|
|
constructor(type, values) {
|
|
this.type = type;
|
|
this.values = values;
|
|
}
|
|
};
|
|
|
|
tnn.TensorType = class {
|
|
|
|
constructor(dataType, shape) {
|
|
this.dataType = dataType || '?';
|
|
this.shape = shape;
|
|
}
|
|
|
|
toString() {
|
|
return this.dataType + this.shape.toString();
|
|
}
|
|
};
|
|
|
|
tnn.TensorShape = class {
|
|
|
|
constructor(dimensions) {
|
|
this.dimensions = dimensions;
|
|
}
|
|
|
|
toString() {
|
|
return this.dimensions ? (`[${this.dimensions.map((dimension) => dimension ? dimension.toString() : '?').join(',')}]`) : '';
|
|
}
|
|
};
|
|
|
|
tnn.TextProtoReader = class {
|
|
|
|
constructor(reader) {
|
|
this.reader = reader;
|
|
this.inputs = [];
|
|
this.outputs = [];
|
|
this.layers = [];
|
|
}
|
|
|
|
read() {
|
|
if (this.reader) {
|
|
let lines = [];
|
|
for (;;) {
|
|
const line = this.reader.read('\n');
|
|
if (line === undefined) {
|
|
break;
|
|
}
|
|
lines.push(line.replace(/\r|"/g, ''));
|
|
}
|
|
const split = (line, delimiter, trim, ignore_blank) => {
|
|
return line.split(delimiter).map((v) => trim ? v.trim() : v).filter((v) => !ignore_blank || v);
|
|
};
|
|
lines = split(lines.join(''), ',', true, false);
|
|
if (lines.length <= 5) {
|
|
throw new tnn.Error('Invalid line count.');
|
|
}
|
|
const header = split(lines.shift(), ' ', true, false);
|
|
if (header.length < 3) {
|
|
throw new tnn.Error('Invalid header size.');
|
|
} else if (header.length > 3 && (header[3] !== '4206624770' && header[3] !== '4206624772')) {
|
|
throw new tnn.Error(`Invalid signature '${header[3]}'.`);
|
|
}
|
|
this.inputs = split(lines.shift(), ':', true, false).map((input) => {
|
|
const array = split(input, ' ', true, false);
|
|
const name = array.shift();
|
|
if (header[3] === '4206624772') {
|
|
const shape_size = parseInt(array.shift(), 10);
|
|
const data_type_index = parseInt(array[shape_size], 10);
|
|
return {
|
|
name,
|
|
data_type: ['float32', 'float16', 'int8', 'int32', 'bfloat16'][data_type_index],
|
|
shape: array.slice(0, -1).map((dim) => parseInt(dim, 10)),
|
|
};
|
|
|
|
}
|
|
return {
|
|
name,
|
|
data_type: 'float32',
|
|
shape: array.map((dim) => parseInt(dim, 10))
|
|
};
|
|
});
|
|
lines.shift();
|
|
this.outputs = split(lines.shift(), ' ', true, false).map((output) => {
|
|
return { name: output };
|
|
});
|
|
lines.shift();
|
|
while (lines.length > 0) {
|
|
const line = lines.shift().trim();
|
|
if (line.length > 0) {
|
|
const array = split(line, ' ', true, true);
|
|
const layer = {};
|
|
layer.type = array.shift();
|
|
layer.name = array.shift();
|
|
const inputs = parseInt(array.shift(), 10);
|
|
const outputs = parseInt(array.shift(), 10);
|
|
layer.inputs = array.splice(0, inputs);
|
|
layer.outputs = array.splice(0, outputs);
|
|
layer.params = new Map();
|
|
let count = 0;
|
|
for (const column of array) {
|
|
const parts = column.split(' ');
|
|
if (parts.length === 1) {
|
|
let key = count.toString();
|
|
let value = parts.toString();
|
|
const keyInt = parseInt(key, 10);
|
|
if (keyInt < 0) {
|
|
value = value.split(',').map((v) => v.trim());
|
|
value.shift();
|
|
key = (-(keyInt + 23300)).toString();
|
|
}
|
|
layer.params.set(key, value);
|
|
count++;
|
|
}
|
|
}
|
|
this.layers.push(layer);
|
|
}
|
|
}
|
|
delete this.reader;
|
|
}
|
|
}
|
|
};
|
|
|
|
tnn.LayerResourceReader = class {
|
|
|
|
static async open(context) {
|
|
if (context) {
|
|
const reader = await context.read('binary');
|
|
return new tnn.LayerResourceReader(reader);
|
|
}
|
|
return new tnn.LayerResourceReader(null);
|
|
}
|
|
|
|
constructor(reader) {
|
|
this.resources = new Map();
|
|
if (reader) {
|
|
this.reader = reader;
|
|
const magic_number = this.reader.uint32();
|
|
if (magic_number !== 0xFABC0002 && magic_number !== 0xFABC0004) {
|
|
throw new tnn.Error(`Invalid blob header signature '${magic_number}'.`);
|
|
}
|
|
const size = this.reader.int32() & 0x1FFFFFFF;
|
|
for (let i = 0; i < size; i++) {
|
|
const resource = {};
|
|
resource.operator = this.reader.int32();
|
|
resource.type = this.reader.string();
|
|
resource.name = this.reader.string();
|
|
switch (resource.type) {
|
|
case 'Convolution':
|
|
case 'ConvolutionDepthWise':
|
|
case 'Deconvolution':
|
|
case 'DeconvolutionDepthWise': {
|
|
this._expect(resource.name);
|
|
const bias = this.reader.int32();
|
|
resource.filter = this._read();
|
|
if (bias) {
|
|
resource.bias = this._read();
|
|
}
|
|
if (resource.filter.dataType === 'int8') {
|
|
resource.quantized = this._read();
|
|
}
|
|
break;
|
|
}
|
|
case 'Conv3D': {
|
|
this._expect(resource.name);
|
|
const bias = this.reader.int32();
|
|
resource.filter = this._read();
|
|
if (bias) {
|
|
resource.bias = this._read();
|
|
}
|
|
break;
|
|
}
|
|
case 'InnerProduct': {
|
|
this._expect(resource.name);
|
|
resource.weight = this._read();
|
|
resource.bias = this._read();
|
|
if (resource.weight.dataType === 'int8') {
|
|
resource.scale = this._read();
|
|
}
|
|
break;
|
|
}
|
|
case 'PReLU': {
|
|
this._expect(resource.name);
|
|
resource.slope = this._read();
|
|
break;
|
|
}
|
|
case 'Add':
|
|
case 'Div':
|
|
case 'Mul':
|
|
case 'Sub':
|
|
case 'MatMul': {
|
|
resource.slope = this._read();
|
|
break;
|
|
}
|
|
case 'BatchNormCxx':
|
|
case 'InstBatchNormCxx':
|
|
resource.scale = this._read();
|
|
resource.bias = this._read();
|
|
break;
|
|
case 'HdrGuide':
|
|
resource.ccm_weight = this._read();
|
|
resource.ccm_bias = this._read();
|
|
resource.shifts = this._read();
|
|
resource.slopes = this._read();
|
|
resource.projection_weight = this._read();
|
|
resource.projection_bias = this._read();
|
|
break;
|
|
case 'BlobScale':
|
|
resource.scale = this._read();
|
|
resource.bias = this._read();
|
|
break;
|
|
case 'Gather': {
|
|
// reader.expect(resource.name);
|
|
const has_data = this.reader.int32();
|
|
if (has_data) {
|
|
resource.data = this._read();
|
|
}
|
|
const has_indices = this.reader.int32();
|
|
if (has_indices) {
|
|
resource.indices = this._read();
|
|
}
|
|
break;
|
|
}
|
|
default: {
|
|
throw new tnn.Error(`Unsupported layer resource type '${resource.type}'.`);
|
|
}
|
|
}
|
|
this.resources.set(resource.name, resource);
|
|
}
|
|
if (this.reader.position !== this.reader.length) {
|
|
throw new tnn.Error("Invalid blob size.");
|
|
}
|
|
delete this.reader;
|
|
}
|
|
}
|
|
|
|
_read() {
|
|
const magic_number = this.reader.uint32();
|
|
if (magic_number !== 0xFABC0002 && magic_number !== 0xFABC0004) {
|
|
throw new tnn.Error(`Invalid raw signature '${magic_number}'.`);
|
|
}
|
|
const data_type = this.reader.int32();
|
|
if (data_type > 4) {
|
|
throw new tnn.Error(`Unsupported data type '${data_type}'.`);
|
|
}
|
|
const length = this.reader.int32();
|
|
if (length <= 0) {
|
|
return null;
|
|
}
|
|
let dims = null;
|
|
if (magic_number === 0xFABC0004) {
|
|
const dim_size = this.reader.int32();
|
|
dims = this.reader.read(dim_size * 4);
|
|
}
|
|
return {
|
|
dataType: ['float32', 'float16', 'int8', 'int32', 'bfloat16'][data_type],
|
|
length: length / [4, 2, 1, 4, 2][data_type],
|
|
value: this.reader.read(length),
|
|
shape: dims
|
|
};
|
|
}
|
|
|
|
_expect(name) {
|
|
const content = this.reader.string();
|
|
if (name !== content) {
|
|
throw new tnn.Error(`Invalid string '${content}' instead of '${name}'.`);
|
|
}
|
|
}
|
|
|
|
get(name) {
|
|
if (this.resources.size === 0) {
|
|
return null;
|
|
}
|
|
if (!this.resources.has(name)) {
|
|
throw new tnn.Error(`Invalid blob layer name '${name}'.`);
|
|
}
|
|
return this.resources.get(name);
|
|
}
|
|
};
|
|
|
|
tnn.Error = class extends Error {
|
|
|
|
constructor(message) {
|
|
super(message);
|
|
this.name = 'Error loading TNN model.';
|
|
}
|
|
};
|
|
|
|
export const ModelFactory = tnn.ModelFactory;
|