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
494 lines
20 KiB
JavaScript
494 lines
20 KiB
JavaScript
|
|
const caffe2 = {};
|
|
|
|
caffe2.ModelFactory = class {
|
|
|
|
async match(context) {
|
|
const identifier = context.identifier.toLowerCase();
|
|
const extension = identifier.lastIndexOf('.') > 0 ? identifier.split('.').pop().toLowerCase() : '';
|
|
switch (extension) {
|
|
case 'pbtxt':
|
|
case 'prototxt': {
|
|
const tags = await context.tags('pbtxt');
|
|
if (tags.has('op') && !tags.has('op.attr') && !tags.has('op.graph_op_name') && !tags.has('op.endpoint')) {
|
|
return context.set('caffe2.pbtxt');
|
|
}
|
|
break;
|
|
}
|
|
case 'pb': {
|
|
const tags = await context.tags('pb');
|
|
if (tags.size > 0 &&
|
|
Array.from(tags.keys()).every((tag) => tag <= 9) &&
|
|
Array.from(tags.values()).every((type) => type <= 4)) {
|
|
if (tags.size === 1 && tags.get(2) === 2 && identifier.endsWith('saved_model.pb')) {
|
|
return null;
|
|
}
|
|
const schema = [[1,2],[2,2],[3,2],[4,0],[5,2],[6,2],[7,2],[8,2],[9,2]];
|
|
if (schema.every(([key, value]) => !tags.has(key) || tags.get(key) === value)) {
|
|
const stream = context.stream;
|
|
if (stream.length > 3) {
|
|
const buffer = stream.peek(Math.min(stream.length, 67));
|
|
const [signature, size] = buffer;
|
|
switch (signature) {
|
|
case 0x0A:
|
|
if (size < 64 &&
|
|
buffer.length > 2 + size + 1 &&
|
|
buffer.slice(2, 2 + size).every((c) => c >= 32 && c <= 127) &&
|
|
buffer[2 + size] === 0x12) {
|
|
return context.set('caffe2.pb');
|
|
}
|
|
break;
|
|
case 0x12:
|
|
return context.set('caffe2.pb');
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
default: {
|
|
break;
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
async open(context) {
|
|
caffe2.proto = await context.require('./caffe2-proto');
|
|
caffe2.proto = caffe2.proto.caffe2;
|
|
const metadata = await context.metadata('caffe2-metadata.json');
|
|
const identifier = context.identifier;
|
|
const parts = identifier.split('.');
|
|
const extension = parts.pop().toLowerCase();
|
|
const base = parts.join('.');
|
|
let predict = null;
|
|
let init = null;
|
|
switch (context.type) {
|
|
case 'caffe2.pbtxt': {
|
|
if (base.toLowerCase().endsWith('init_net') || base.toLowerCase().startsWith('init_net')) {
|
|
init = context;
|
|
try {
|
|
const name = identifier.replace('init_net', 'predict_net');
|
|
predict = await context.fetch(name);
|
|
predict.set(context.type);
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
} else if (base.toLowerCase().endsWith('predict_net') || base.toLowerCase().startsWith('predict_net')) {
|
|
predict = context;
|
|
const name = identifier.replace('predict_net', 'init_net');
|
|
try {
|
|
init = await context.fetch(name.replace(/\.pbtxt/, '.pb'));
|
|
init.set('caffe2.pb');
|
|
} catch {
|
|
try {
|
|
init = await context.fetch(name);
|
|
init.set('caffe2.pbtxt');
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
}
|
|
} else {
|
|
predict = context;
|
|
try {
|
|
init = await context.fetch(`${base}_init.pb`);
|
|
init.set('caffe2.pb');
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case 'caffe2.pb': {
|
|
if (base.toLowerCase().endsWith('init_net')) {
|
|
init = context;
|
|
const extensions = new Set([extension, 'pb', 'pbtxt']);
|
|
for (const extension of extensions) {
|
|
try {
|
|
const name = `${base.replace(/init_net$/, '')}predict_net.${extension}`;
|
|
// eslint-disable-next-line no-await-in-loop
|
|
predict = await context.fetch(name);
|
|
predict.set(`caffe2.${extension}`);
|
|
break;
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
}
|
|
} else if (base.toLowerCase().endsWith('_init')) {
|
|
try {
|
|
const name = `${base.replace(/_init$/, '')}.${extension}`;
|
|
predict = await context.fetch(name);
|
|
predict.set(context.type);
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
} else if (base.toLowerCase().endsWith('predict_net') || base.toLowerCase().startsWith('predict_net')) {
|
|
predict = context;
|
|
try {
|
|
const name = identifier.replace('predict_net', 'init_net');
|
|
init = await context.fetch(name);
|
|
init.set(context.type);
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
} else {
|
|
predict = context;
|
|
try {
|
|
const file = `${base}_init.${extension}`;
|
|
init = await context.fetch(file, null);
|
|
init.set(context.type);
|
|
} catch {
|
|
// continue regardless of error
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
default: {
|
|
throw new caffe2.Error(`Unsupported Caffe2 format '${context.type}'.`);
|
|
}
|
|
}
|
|
if (!predict && !init) {
|
|
throw new caffe2.Error(`Caffe2 model does not contain predict or init data.`);
|
|
}
|
|
const open = async (context) => {
|
|
if (context) {
|
|
switch (context.type) {
|
|
case 'caffe2.pb':
|
|
try {
|
|
const reader = await context.read('protobuf.binary');
|
|
return caffe2.proto.NetDef.decode(reader);
|
|
} catch (error) {
|
|
const message = error && error.message ? error.message : error.toString();
|
|
throw new caffe2.Error(`File format is not caffe2.NetDef (${message.replace(/\.$/, '')}).`);
|
|
}
|
|
case 'caffe2.pbtxt':
|
|
try {
|
|
const reader = await context.read('protobuf.text');
|
|
reader.field = function(tag, message) {
|
|
if (message instanceof caffe2.proto.DeviceOption) {
|
|
message[tag] = this.read();
|
|
return;
|
|
}
|
|
throw new Error(`Unknown field '${tag}' ${this.location()}`);
|
|
};
|
|
return caffe2.proto.NetDef.decodeText(reader);
|
|
} catch (error) {
|
|
const message = error && error.message ? error.message : error.toString();
|
|
throw new caffe2.Error(`File format is not caffe2.NetDef (${message.replace(/\.$/, '')}).`);
|
|
}
|
|
default:
|
|
throw new caffe2.Error(`Unsupported Caffe2 predict format '${context.type}'.`);
|
|
}
|
|
}
|
|
return null;
|
|
};
|
|
const predict_net = await open(predict);
|
|
const init_net = await open(init);
|
|
return new caffe2.Model(metadata, predict_net, init_net);
|
|
}
|
|
};
|
|
|
|
caffe2.Model = class {
|
|
|
|
constructor(metadata, predict_net, init_net) {
|
|
const net = predict_net || init_net;
|
|
this.format = 'Caffe2';
|
|
this.domain = net.domain || null;
|
|
this.modules = [new caffe2.Graph(metadata, predict_net, init_net)];
|
|
}
|
|
};
|
|
|
|
caffe2.Graph = class {
|
|
|
|
constructor(metadata, predict_net, init_net) {
|
|
const net = predict_net || init_net;
|
|
init_net = predict_net ? init_net : null;
|
|
this.name = net.name || '';
|
|
this.nodes = [];
|
|
this.description = net.type;
|
|
const initializers = new Set();
|
|
const tensors = new Map();
|
|
for (const name of net.external_input) {
|
|
tensors.set(name, new caffe2.Tensor(name, {}));
|
|
}
|
|
if (init_net) {
|
|
const dataTypes = new Map([
|
|
['GivenTensorFill', 'float32'],
|
|
['GivenTensorDoubleFill', 'float64'],
|
|
['GivenTensorBoolFill', 'boolean'],
|
|
['GivenTensorByteStringToUInt8Fill', 'uint8'],
|
|
['GivenTensorInt16Fill', 'int16'],
|
|
['GivenTensorSInt16Fill', 'int16'],
|
|
['GivenTensorIntFill', 'int32'],
|
|
['GivenTensorInt64Fill', 'int64'],
|
|
['GivenTensorStringFill', 'string'],
|
|
['Int8GivenIntTensorFill', 'int32'],
|
|
['Int8GivenTensorFill', 'int8'],
|
|
['XavierFill', null],
|
|
['ConstantFill', null]
|
|
]);
|
|
for (const op of init_net.op) {
|
|
if (op.output && op.output.length === 1) {
|
|
const [name] = op.output;
|
|
const tensor = {};
|
|
for (const arg of op.arg) {
|
|
tensor[arg.name] = arg;
|
|
}
|
|
if (!dataTypes.has(op.type)) {
|
|
throw new caffe2.Error(`Unsupported init op '${op.type}'.`);
|
|
}
|
|
tensor.dataType = dataTypes.get(op.type);
|
|
if (tensor.values && tensor.values.floats && (tensor.values.floats.length !== 1 || tensor.values.floats[0] !== 0)) {
|
|
initializers.add(name);
|
|
}
|
|
tensors.set(name, new caffe2.Tensor(name, tensor));
|
|
}
|
|
}
|
|
}
|
|
const scope = {};
|
|
for (let i = 0; i < net.op.length; i++) {
|
|
const op = net.op[i];
|
|
op.input = op.input.map((input) => scope[input] ? scope[input] : input);
|
|
op.output = op.output.map((output) => {
|
|
if (scope[output]) {
|
|
const next = `${output}\n${i}`; // custom argument id
|
|
scope[output] = next;
|
|
return next;
|
|
}
|
|
scope[output] = output;
|
|
return output;
|
|
});
|
|
}
|
|
const values = new Map();
|
|
values.map = (name, type, tensor) => {
|
|
if (!values.has(name)) {
|
|
values.set(name, new caffe2.Value(name, type || null, tensor || null));
|
|
} else if (type || tensor) {
|
|
throw new caffe2.Value(`Duplicate value '${name}'.`);
|
|
}
|
|
return values.get(name);
|
|
};
|
|
for (const op of net.op) {
|
|
let index = 0;
|
|
for (const name of op.input) {
|
|
if (index > 0 && tensors.has(name)) {
|
|
if (!values.has(name)) {
|
|
values.set(name, new caffe2.Value(name, null, tensors.get(name)));
|
|
}
|
|
initializers.add(name);
|
|
}
|
|
index++;
|
|
}
|
|
}
|
|
for (const op of net.op) {
|
|
for (const name of op.output) {
|
|
if (tensors.has(name)) {
|
|
initializers.add(name);
|
|
}
|
|
}
|
|
}
|
|
let lastNode = null;
|
|
let lastOutput = null;
|
|
for (const op of net.op) {
|
|
const node = new caffe2.Node(metadata, op, values);
|
|
if (op.input.length === 1 &&
|
|
op.output.length >= 1 &&
|
|
op.input[0].split('\n').shift() === op.output[0].split('\n').shift() &&
|
|
lastNode &&
|
|
lastOutput === op.input[0].split('\n').shift()) {
|
|
lastNode.chain.push(node);
|
|
} else {
|
|
this.nodes.push(node);
|
|
lastNode = null;
|
|
lastOutput = null;
|
|
if (op.output.length === 1) {
|
|
lastNode = node;
|
|
lastOutput = op.output[0].split('\n').shift();
|
|
}
|
|
}
|
|
}
|
|
this.inputs = [];
|
|
for (const input of net.external_input) {
|
|
if (net.external_input.length > 1 && initializers.has(input)) {
|
|
continue;
|
|
}
|
|
const argument = new caffe2.Argument(input, [values.map(input)]);
|
|
this.inputs.push(argument);
|
|
}
|
|
this.outputs = [];
|
|
for (const output of net.external_output) {
|
|
const argument = new caffe2.Argument(output, [values.map(output)]);
|
|
this.outputs.push(argument);
|
|
}
|
|
}
|
|
};
|
|
|
|
caffe2.Argument = class {
|
|
|
|
constructor(name, value, type = null, visible = true) {
|
|
this.name = name;
|
|
this.value = value;
|
|
this.type = type;
|
|
this.visible = visible;
|
|
}
|
|
};
|
|
|
|
caffe2.Value = class {
|
|
|
|
constructor(name, type, initializer = null) {
|
|
if (typeof name !== 'string') {
|
|
throw new caffe2.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
|
|
}
|
|
this.name = name;
|
|
this.type = !type && initializer ? initializer.type : type;
|
|
this.quantization = initializer && initializer.quantization ? initializer.quantization : null;
|
|
this.initializer = initializer;
|
|
}
|
|
};
|
|
|
|
caffe2.Node = class {
|
|
|
|
constructor(metadata, op, values) {
|
|
this.name = op.name || '';
|
|
this.device = op.engine || '';
|
|
this.chain = [];
|
|
this.type = metadata.type(op.type);
|
|
this.attributes = op.arg.map((arg) => {
|
|
const schema = metadata.attribute(op.type, arg.name);
|
|
const name = arg.name;
|
|
let value = null;
|
|
let type = null;
|
|
let visible = true;
|
|
if (arg.floats && arg.floats.length > 0) {
|
|
value = arg.floats;
|
|
} else if (arg.ints && arg.ints.length > 0) {
|
|
value = arg.ints;
|
|
} else if (arg.nets && arg.nets.length > 0) {
|
|
value = arg.nets.map((net) => new caffe2.Graph(metadata, net, null));
|
|
type = 'graph[]';
|
|
} else if (arg.n) {
|
|
value = new caffe2.Graph(metadata, arg.n, null);
|
|
type = 'graph';
|
|
} else {
|
|
value = arg.i;
|
|
}
|
|
if (schema) {
|
|
type = !type && schema.type ? schema.type : type;
|
|
if (type === 'boolean') {
|
|
value = value !== 0 && value.toString() !== '0' ? true : false;
|
|
}
|
|
if (schema.visible === false) {
|
|
visible = false;
|
|
} else if (schema.default !== undefined) {
|
|
if (value === metadata.default || (value && value.toString() === schema.default.toString())) {
|
|
visible = false;
|
|
}
|
|
}
|
|
}
|
|
return new caffe2.Argument(name, value, type, visible);
|
|
});
|
|
const inputs = op.input;
|
|
const outputs = op.output;
|
|
this.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.option === 'variadic') ? (inputs.length - inputIndex) : 1;
|
|
const inputArguments = inputs.slice(inputIndex, inputIndex + inputCount).filter((id) => id !== '' || inputDef.option !== 'optional').map((id) => values.map(id));
|
|
this.inputs.push(new caffe2.Argument(inputDef.name, inputArguments));
|
|
inputIndex += inputCount;
|
|
}
|
|
}
|
|
} else {
|
|
this.inputs.push(...inputs.slice(inputIndex).map((input, index) => {
|
|
const inputName = ((inputIndex + index) === 0) ? 'input' : (inputIndex + index).toString();
|
|
return new caffe2.Argument(inputName, [values.map(input)]);
|
|
}));
|
|
}
|
|
this.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));
|
|
this.outputs.push(new caffe2.Argument(outputDef.name, outputArguments));
|
|
outputIndex += outputCount;
|
|
}
|
|
}
|
|
} else {
|
|
this.outputs.push(...outputs.slice(outputIndex).map((output, index) => {
|
|
const outputName = ((outputIndex + index) === 0) ? 'output' : (outputIndex + index).toString();
|
|
return new caffe2.Argument(outputName, [values.map(output)]);
|
|
}));
|
|
}
|
|
}
|
|
};
|
|
|
|
caffe2.Tensor = class {
|
|
|
|
constructor(name, tensor) {
|
|
this.name = name;
|
|
const shape = tensor.shape && tensor.shape.ints ? tensor.shape.ints : null;
|
|
this.type = new caffe2.TensorType(tensor.dataType, new caffe2.TensorShape(shape));
|
|
this.values = null;
|
|
this.category = 'Initializer';
|
|
this.encoding = '|';
|
|
if (tensor.Y_scale !== undefined || tensor.Y_zero_point !== undefined) {
|
|
this.quantization = {
|
|
type: 'linear',
|
|
scale: [tensor.Y_scale ? tensor.Y_scale.f : 0],
|
|
offset: [tensor.Y_zero_point && typeof tensor.Y_zero_point.i === 'bigint' ? tensor.Y_zero_point.i.toNumber() : 0]
|
|
};
|
|
}
|
|
if (tensor.values) {
|
|
switch (this.type.dataType) {
|
|
case 'float32': this.values = tensor.values.floats; break;
|
|
case 'boolean': this.values = tensor.values.ints; break;
|
|
case 'int8': this.values = new Int8Array(tensor.values.s); break;
|
|
case 'int32': this.values = tensor.values.ints; break;
|
|
default: break;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
caffe2.TensorType = class {
|
|
|
|
constructor(dataType, shape) {
|
|
this.dataType = dataType || '?';
|
|
this.shape = shape;
|
|
}
|
|
|
|
toString() {
|
|
return this.dataType + this.shape.toString();
|
|
}
|
|
};
|
|
|
|
caffe2.TensorShape = class {
|
|
|
|
constructor(dimensions) {
|
|
this.dimensions = Array.isArray(dimensions) ? dimensions.map((dim) => typeof dim === 'bigint' ? dim.toNumber() : dim) : dimensions;
|
|
}
|
|
|
|
toString() {
|
|
if (Array.isArray(this.dimensions) && this.dimensions.length > 0) {
|
|
return `[${this.dimensions.map((dim) => dim.toString()).join(',')}]`;
|
|
}
|
|
return '';
|
|
}
|
|
};
|
|
|
|
caffe2.Error = class extends Error {
|
|
|
|
constructor(message) {
|
|
super(message);
|
|
this.name = 'Error loading Caffe2 model.';
|
|
}
|
|
};
|
|
|
|
export const ModelFactory = caffe2.ModelFactory;
|