// Experimental import * as base from './base.js'; import * as protobuf from './protobuf.js'; import * as zip from './zip.js'; const tf = {}; tf.ModelFactory = class { async match(context) { const identifier = context.identifier; const extension = identifier.lastIndexOf('.') > 0 ? identifier.split('.').pop().toLowerCase() : ''; if (extension === 'pbtxt' || extension === 'prototxt' || extension === 'pt' || extension === 'txt') { if (identifier.endsWith('predict_net.pbtxt') || identifier.endsWith('predict_net.prototxt') || identifier.endsWith('init_net.pbtxt') || identifier.endsWith('init_net.prototxt')) { return null; } const tags = await context.tags('pbtxt'); if (['input_stream', 'output_stream', 'input_side_packet', 'output_side_packet'].some((key) => tags.has(key) || tags.has(`node.${key}`))) { return null; } if (tags.has('saved_model_schema_version') || tags.has('meta_graphs')) { return context.set('tf.pbtxt.SavedModel'); } if (tags.has('graph_def')) { return context.set('tf.pbtxt.MetaGraphDef'); } if (tags.has('node')) { return context.set('tf.pbtxt.GraphDef'); } } if (extension === 'pb' || extension === 'pbtxt' || extension === 'prototxt' || extension === 'graphdef' || extension === 'meta') { if (identifier.endsWith('predict_net.pb') || identifier.endsWith('init_net.pb')) { return null; } if (identifier === 'tfhub_module.pb') { const stream = context.stream; const signature = [0x08, 0x03]; if (signature.length === stream.length && stream.peek(signature.length).every((value, index) => value === signature[index])) { return null; } } const tags = await context.tags('pb'); if (tags.size > 0) { if (Array.from(tags).every(([key, value]) => key < 8 && value !== 5)) { const match = (tags, schema) => { for (const [key, inner] of schema) { const value = tags[key]; if (value === undefined) { continue; } if (inner === false) { return false; } if (Array.isArray(inner)) { if (typeof value !== 'object' || !match(value, inner)) { return false; } } else if (inner !== value) { if (inner === 2 && !Array.isArray(value) && Object(value) === (value) && Object.keys(value).length === 0) { return true; } return false; } } return true; }; const signatureGraphDef = [ [1 /* node */, [ [1 /* name */, 2], [2 /* op */, 2], [3 /* input */, 2], [4 /* device */,2], [5 /* attr */, [ [1,2], [2,[]] ]], [6 /* experimental_debug_info */, []] ]], [2 /* library */, []], [3 /* version */, 0], [4 /* versions */, [[1,0],[2,0]]] ]; const signatureMetaGraphDef = [ [1 /* meta_info_def */, [[1,2],[2,[]],[3,[]],/* [4,2], */[6,2],[7,0],[8,[]]]], [2 /* graph_def */, signatureGraphDef], [3 /* saver_def */, [[1,2],[2,2],[3,2],[4,0],[5,0],[6,5],[7,0]]], [4 /* collection_def */,[]], [5 /* signature_def */, []], [6 /* asset_file_def */, []], [7 /* object_graph_def */, []] ]; const signatureSavedModel = [[1,0],[2,signatureMetaGraphDef]]; // optimization_guide.proto.PageTopicsOverrideList if (identifier === 'override_list.pb' && tags.size === 1 && tags.get(1) === 2) { return null; } if (tags.size === 1 && tags.get(1) === 2) { const tags = await context.tags('pb+'); // mediapipe.BoxDetectorIndex if (match(tags, [[1,[[1,[[1,[[1,5],[2,5],[3,5],[4,5],[6,0],[7,5],[8,5],[10,5],[11,0],[12,0]]],[2,5],[3,[]]]],[2,false],[3,false],[4,false],[5,false]]],[2,false],[3,false]])) { return null; } // third_party.tensorflow.python.keras.protobuf.SavedMetadata if (match(tags, [[1,[[1,[[1,0],[2,0]]],[2,0],[3,2],[4,2],[5,2]]]])) { return null; } } if ((!tags.has(1) || tags.get(1) === 0) && tags.get(2) === 2) { const tags = await context.tags('pb+'); if (match(tags, signatureSavedModel)) { return context.set('tf.pb.SavedModel'); } } if ((!tags.has(1) || tags.get(1) === 2) && (!tags.has(2) || tags.get(2) === 2) && (!tags.has(3) || tags.get(3) === 2) && (!tags.has(4) || tags.get(4) === 2)) { const tags = await context.tags('pb+'); if (match(tags, signatureMetaGraphDef)) { return context.set('tf.pb.MetaGraphDef'); } } if (tags.get(1) !== 2) { const tags = await context.tags('pb+'); if (match(tags, signatureGraphDef)) { return context.set('tf.pb.GraphDef'); } } // tensorflow.FingerprintDef if (identifier === 'fingerprint.pb' && tags.get(1) === 0 && tags.get(2) === 0 && tags.get(3) === 0 && tags.get(5) === 0 && tags.get(6) === 2) { return context.set('tf.pb.FingerprintDef'); } const decode = (buffer, value) => { try { const reader = protobuf.BinaryReader.open(buffer); const length = reader.length; while (reader.position < length) { const tag = reader.uint32(); const number = tag >>> 3; const type = tag & 7; if (value === number) { return type === 2 ? reader.bytes() : null; } reader.skipType(type); } } catch { // continue regardless of error } return null; }; const stream = context.stream; const buffer = stream.peek(); const nodeBuffer = decode(buffer, 1); if (nodeBuffer) { const nameBuffer = decode(nodeBuffer, 1); if (nameBuffer) { const decoder = new TextDecoder('utf-8'); const name = decoder.decode(nameBuffer); if (Array.from(name).filter((c) => c <= ' ').length < 256) { return context.set('tf.pb.GraphDef'); } } } } } else { const tags = await context.tags('pbtxt'); if (['input_stream', 'output_stream', 'input_side_packet', 'output_side_packet'].some((key) => tags.has(key) || tags.has(`node.${key}`))) { return null; } if (tags.has('node')) { return context.set('tf.pbtxt.GraphDef'); } if (tags.has('graph_def')) { return context.set('tf.pbtxt.MetaGraphDef'); } if (tags.has('saved_model_schema_version') || tags.has('meta_graphs')) { return context.set('tf.pbtxt.SavedModel'); } } } if (extension === 'json') { for (const type of ['json', 'json.gz']) { // eslint-disable-next-line no-await-in-loop const obj = await context.peek(type); if (obj && obj.modelTopology && (obj.format === 'graph-model' || Array.isArray(obj.modelTopology.node))) { return context.set(`tf.${type}`); } } } if (extension === 'index' || extension === 'ckpt') { const stream = context.stream; if (stream.length > 8) { stream.seek(-8); const buffer = stream.read(8); stream.seek(0); const signature = [0x57, 0xfb, 0x80, 0x8b, 0x24, 0x75, 0x47, 0xdb]; if (buffer.every((value, index) => value === signature[index])) { return context.set('tf.bundle'); } } } if (/.data-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]$/.exec(identifier)) { return context.set('tf.data'); } if (/^events.out.tfevents./.exec(identifier)) { const stream = context.stream; if (tf.EventFileReader.open(stream)) { return context.set('tf.events'); } } if (extension === 'pbmm') { const stream = context.stream; if (stream.length > 8) { stream.seek(-8); const buffer = stream.read(8); stream.seek(0); const reader = base.BinaryReader.open(buffer); const offset = reader.uint64().toNumber(); if (offset < stream.length) { return context.set('tf.pb.mmap'); } } } if (/^.*group\d+-shard\d+of\d+(\.bin)?$/.test(identifier)) { return context.set('tf.tfjs.weights'); } return null; } filter(context, match) { if (context.type === 'tf.bundle' && match.type === 'tf.data') { return false; } if ((context.type === 'tf.json' || context.type === 'tf.json.gz') && match.type === 'tf.tfjs.weights') { return false; } return true; } async open(context) { tf.proto = await context.require('./tf-proto'); const openModel = async (saved_model, format, producer, bundle) => { const metadata = await context.metadata('tf-metadata.json'); return new tf.Model(metadata, saved_model, format, producer, bundle); }; const openSavedModel = async (context, saved_model, format, producer) => { if (format === '') { format = 'TensorFlow Saved Model'; if (saved_model && saved_model.saved_model_schema_version) { format = `${format} v${saved_model.saved_model_schema_version}`; } } if (saved_model.meta_graphs.length === 1 && saved_model.meta_graphs[0].object_graph_def && saved_model.meta_graphs[0].object_graph_def.nodes && saved_model.meta_graphs[0].object_graph_def.nodes.length > 0) { const identifier = 'variables/variables.index'; try { const content = await context.fetch(identifier); const stream = content.stream; const bundle = await tf.TensorBundle.open(stream, identifier, context); return openModel(saved_model, format, producer, bundle); } catch { return openModel(saved_model, format, producer, null); } } if (saved_model && Array.isArray(saved_model.meta_graphs) && saved_model.meta_graphs.length > 0 && saved_model.meta_graphs[0].meta_info_def && Object.prototype.hasOwnProperty.call(saved_model.meta_graphs[0].meta_info_def, 'tensorflow_version')) { producer = `TensorFlow v${saved_model.meta_graphs[0].meta_info_def.tensorflow_version}`; } return openModel(saved_model, format, producer, null); }; const openBundle = async (context, stream, identifier) => { stream = stream || context.stream; identifier = identifier || context.identifier; try { const bundle = await tf.TensorBundle.open(stream, identifier, context); return openModel(null, `TensorFlow Tensor Bundle v${bundle.format}`, null, bundle); } catch (error) { context.error(error, false); throw error; } }; const openData = async (context) => { const identifier = context.identifier; const base = identifier.split('.'); base.pop(); const file = `${base.join('.')}.index`; try { const content = await context.fetch(file); const stream = content.stream; return openBundle(context, stream, file); } catch { const file = `${base.join('.')}.ckpt`; const content = await context.fetch(file); const stream = content.stream; return openBundle(context, stream, file); } }; const openEventFile = async (context) => { let format = 'TensorFlow Event File'; let producer = null; const stream = context.stream; const eventFileReader = tf.EventFileReader.open(stream); const saved_model = new tf.proto.tensorflow.SavedModel(); const run_metadata = []; const summaries = []; for (;;) { const event = eventFileReader.read(); if (!event) { break; } switch (event.what) { case 'file_version': { const formats = new Map([ ['brain.Event:1', 'TensorFlow Event File v1'], ['brain.Event:2', 'TensorFlow Event File v2'] ]); if (!formats.has(event.file_version)) { throw new tf.Error(`Unsupported event file version '${event.file_version}'.`); } format = formats.get(event.file_version); break; } case 'graph_def': { const buffer = event.graph_def; const reader = protobuf.BinaryReader.open(buffer); const graph_def = tf.proto.tensorflow.GraphDef.decode(reader); const meta_graph_def = new tf.proto.tensorflow.MetaGraphDef(); meta_graph_def.meta_info_def = new tf.proto.tensorflow.MetaGraphDef.MetaInfoDef(); meta_graph_def.meta_info_def.any_info = event.wall_time.toString(); meta_graph_def.graph_def = graph_def; saved_model.meta_graphs.push(meta_graph_def); break; } case 'meta_graph_def': { const buffer = event.meta_graph_def; const reader = protobuf.BinaryReader.open(buffer); const meta_graph_def = tf.proto.tensorflow.MetaGraphDef.decode(reader); saved_model.meta_graphs.push(meta_graph_def); break; } case 'summary': { for (const value of event.summary.value) { summaries.push(value); } break; } case 'tagged_run_metadata': { const entry = event.tagged_run_metadata; const buffer = entry.run_metadata; const reader = protobuf.BinaryReader.open(buffer); const metadata = tf.proto.tensorflow.RunMetadata.decode(reader); run_metadata.push(metadata); break; } default: { throw new tf.Error(`Unsupported event type '${event.what}'.`); } } } if (saved_model.meta_graphs.every((meta_graph) => meta_graph.graph_def.node.every((node) => node.op.startsWith('aten::') || node.op.startsWith('prim::') || node.op.startsWith('quantized::') || node.op === 'IO Node'))) { producer = 'PyTorch'; const openPyTorchMetadata = async (context, saved_model) => { try { const pytorch = await context.require('./pytorch'); const python = await context.require('./python'); const metadata = await pytorch.Metadata.open(context); const execution = new python.Execution(); metadata.register(execution); const torch = execution.__import__('torch'); for (const graph of saved_model.meta_graphs) { for (const node of graph.graph_def.node) { const schemas = torch._C._jit_get_schemas_for_operator(node.op); if (Array.isArray(schemas) && schemas.length > 0) { node.__metadata__ = schemas; node.__torch__ = torch; } } } } catch { // continue regardless of error } return saved_model; }; const updated_saved_model = await openPyTorchMetadata(context, saved_model); return await openModel(updated_saved_model, format, producer, null); } return await openSavedModel(context, saved_model, format, producer); }; const openJson = async (context, type) => { const obj = await context.peek(type); if (!obj || !obj.modelTopology || (obj.format !== 'graph-model' && !Array.isArray(obj.modelTopology.node))) { throw new tf.Error('File format is not TensorFlow.js graph-model.'); } const format = `TensorFlow.js ${obj.format || 'graph-model'}`; const producer = obj.convertedBy || obj.generatedBy || ''; const meta_graph = new tf.proto.tensorflow.MetaGraphDef(); meta_graph.graph_def = tf.proto.tensorflow.GraphDef.decodeJson(obj.modelTopology); const saved_model = new tf.proto.tensorflow.SavedModel(); saved_model.meta_graphs.push(meta_graph); const nodes = new Map(); for (const node of meta_graph.graph_def.node) { node.input = node.input || []; if (node.op === 'Const') { nodes.set(node.name, node); } } const shards = new Map(); const manifests = Array.isArray(obj.weightsManifest) ? obj.weightsManifest : []; for (const manifest of manifests) { for (const path of manifest.paths) { if (!shards.has(path)) { shards.set(path, context.fetch(path)); } } } const openShards = (shards) => { const dtype_size_map = new Map([ ['float16', 2], ['float32', 4], ['float64', 8], ['int8', 1], ['int16', 2], ['int32', 4], ['int64', 8], ['uint8', 1], ['uint16', 2], ['uint32', 4], ['uint64', 8], ['bool', 1] ]); for (const manifest of manifests) { let buffer = null; if (Array.isArray(manifest.paths) && manifest.paths.length > 0 && manifest.paths.every((path) => shards.has(path))) { const list = manifest.paths.map((path) => shards.get(path)); const size = list.reduce((a, b) => a + b.length, 0); buffer = new Uint8Array(size); let offset = 0; for (const item of list) { buffer.set(item, offset); offset += item.length; } } let offset = 0; for (const weight of manifest.weights) { const dtype = weight.quantization && weight.quantization.dtype ? weight.quantization.dtype : weight.dtype; const size = weight.shape.reduce((a, b) => a * b, 1); switch (dtype) { case 'string': { const data = []; if (buffer && size > 0) { const reader = new tf.BinaryReader(buffer.subarray(offset)); for (let i = 0; i < size; i++) { data[i] = reader.string(); } offset += reader.position; } if (nodes.has(weight.name)) { const node = nodes.get(weight.name); node.attr.value.tensor.dtype = tf.Utility.dataTypeKey(dtype); node.attr.value.tensor.string_val = data; } break; } default: { if (!dtype_size_map.has(dtype)) { throw new tf.Error(`Unsupported weight data type size '${dtype}'.`); } const itemsize = dtype_size_map.get(dtype); const length = itemsize * size; const tensor_content = buffer ? buffer.slice(offset, offset + length) : null; offset += length; if (nodes.has(weight.name)) { const node = nodes.get(weight.name); node.attr.value.tensor.dtype = tf.Utility.dataTypeKey(dtype); node.attr.value.tensor.tensor_content = tensor_content; } break; } } } } return openSavedModel(context, saved_model, format, producer); }; try { const contexts = await Promise.all(shards.values()); for (const key of shards.keys()) { const context = contexts.shift(); const buffer = context.stream.peek(); shards.set(key, buffer); } if (type === 'json.gz') { try { for (const key of shards.keys()) { const stream = shards.get(key); const archive = zip.Archive.open(stream, 'gzip'); if (archive && archive.entries.size === 1) { const stream = archive.entries.values().next().value; const buffer = stream.peek(); shards.set(key, buffer); } } } catch { // continue regardless of error } } return openShards(shards); } catch { shards.clear(); return openShards(shards); } }; const openJsonWeights = async (context) => { const content = await context.fetch('model.json'); return await openJson(content, 'json'); }; const openTextGraphDef = async (context) => { try { const reader = await context.read('protobuf.text'); const graph_def = tf.proto.tensorflow.GraphDef.decodeText(reader); const meta_graph = new tf.proto.tensorflow.MetaGraphDef(); meta_graph.graph_def = graph_def; const saved_model = new tf.proto.tensorflow.SavedModel(); saved_model.meta_graphs.push(meta_graph); const format = 'TensorFlow Graph'; return openSavedModel(context, saved_model, format, null); } catch (error) { const message = error && error.message ? error.message : error.toString(); throw new tf.Error(`File text format is not tensorflow.GraphDef (${message.replace(/\.$/, '')}).`); } }; const openTextMetaGraphDef = async (context) => { try { const reader = await context.read('protobuf.text'); const meta_graph = tf.proto.tensorflow.MetaGraphDef.decodeText(reader); const saved_model = new tf.proto.tensorflow.SavedModel(); saved_model.meta_graphs.push(meta_graph); const format = 'TensorFlow MetaGraph'; return openSavedModel(context, saved_model, format, null); } catch (error) { throw new tf.Error(`File text format is not tensorflow.MetaGraphDef (${error.message}).`); } }; const openTextSavedModel = async (context) => { try { const reader = await context.read('protobuf.text'); return tf.proto.tensorflow.SavedModel.decodeText(reader); } catch (error) { throw new tf.Error(`File text format is not tensorflow.SavedModel (${error.message}).`); } }; const openBinaryGraphDef = async (context) => { let saved_model = null; const format = 'TensorFlow Graph'; try { const reader = await context.read('protobuf.binary'); const graph_def = tf.proto.tensorflow.GraphDef.decode(reader); const meta_graph = new tf.proto.tensorflow.MetaGraphDef(); meta_graph.graph_def = graph_def; saved_model = new tf.proto.tensorflow.SavedModel(); saved_model.meta_graphs.push(meta_graph); } catch (error) { const message = error && error.message ? error.message : error.toString(); throw new tf.Error(`File format is not tensorflow.GraphDef (${message.replace(/\.$/, '')}).`); } return openSavedModel(context, saved_model, format, null); }; const openBinaryMetaGraphDef = async (context) => { let saved_model = null; const format = 'TensorFlow MetaGraph'; try { const reader = await context.read('protobuf.binary'); const meta_graph = tf.proto.tensorflow.MetaGraphDef.decode(reader); saved_model = new tf.proto.tensorflow.SavedModel(); saved_model.meta_graphs.push(meta_graph); } catch (error) { const message = error && error.message ? error.message : error.toString(); throw new tf.Error(`File format is not tensorflow.MetaGraphDef (${message.replace(/\.$/, '')}).`); } return openSavedModel(context, saved_model, format, null); }; const openBinarySavedModel = async (context) => { try { const reader = await context.read('protobuf.binary'); return tf.proto.tensorflow.SavedModel.decode(reader); } catch (error) { const message = error && error.message ? error.message : error.toString(); throw new tf.Error(`File format is not tensorflow.SavedModel (${message.replace(/\.$/, '')}).`); } }; const openFingerprint = async (context) => { let format = ''; let saved_model = null; try { const identifier = 'saved_model.pb'; const content = await context.fetch(identifier); saved_model = await openBinarySavedModel(content); } catch { format = 'TensorFlow Fingerprint'; saved_model = new tf.proto.tensorflow.SavedModel(); } const reader = await context.read('protobuf.binary'); saved_model.fingerprint = tf.proto.tensorflow.FingerprintDef.decode(reader); return await openSavedModel(context, saved_model, format, null); }; const openMemmapped = async (context) => { const stream = context.stream; const readDirectoryOffset = (stream) => { stream.seek(-8); stream = stream.stream(8); const reader = base.BinaryReader.open(stream); return reader.uint64().toNumber(); }; const readDirectory = (stream, offset) => { const end = stream.position - 8; stream.seek(offset); stream = stream.stream(end - offset); const reader = protobuf.BinaryReader.open(stream); return tf.proto.tensorflow.MemmappedFileSystemDirectory.decode(reader); }; const offset = readDirectoryOffset(stream); const directory = readDirectory(stream, offset); const elements = new Map(); for (const element of directory.element) { const name = element.name; if (elements.has(name)) { throw new tf.Error(`Memory mapped file directory contains duplicate '${name}'.`); } elements.set(name, { offset: typeof element.offset === 'bigint' ? Number(element.offset) : element.offset, length: typeof element.length === 'bigint' ? Number(element.length) : element.length }); } const offsets = Array.from(elements).map(([, value]) => value.offset); offsets.push(offset); for (const value of elements.values()) { if (value.length === 0) { const min = Math.min.apply(null, offsets.filter((offset) => offset > value.offset)); if (Number.isInteger(min)) { value.length = min - value.offset; } } } for (const [, value] of elements) { const offset = value.offset; const length = value.length; stream.seek(offset); value.buffer = stream.read(length); } if (!elements.has('memmapped_package://.')) { throw new tf.Error('Memory mapped file directory does not contain tensorflow.GraphDef root.'); } const element = elements.get('memmapped_package://.'); const buffer = element.buffer; const reader = protobuf.BinaryReader.open(buffer); const graph_def = tf.proto.tensorflow.GraphDef.decode(reader); const format = 'TensorFlow GraphDef Memmapped'; const meta_graph = new tf.proto.tensorflow.MetaGraphDef(); meta_graph.graph_def = graph_def; const saved_model = new tf.proto.tensorflow.SavedModel(); saved_model.meta_graphs.push(meta_graph); return openSavedModel(context, saved_model, format, null); }; switch (context.type) { case 'tf.bundle': return await openBundle(context); case 'tf.data': return await openData(context); case 'tf.events': return await openEventFile(context); case 'tf.json': return await openJson(context, 'json'); case 'tf.json.gz': return await openJson(context, 'json.gz'); case 'tf.tfjs.weights': return await openJsonWeights(context); case 'tf.pbtxt.GraphDef': return await openTextGraphDef(context); case 'tf.pbtxt.MetaGraphDef': return await openTextMetaGraphDef(context); case 'tf.pbtxt.SavedModel': return await openSavedModel(context, await openTextSavedModel(context), '', null); case 'tf.pb.GraphDef': return await openBinaryGraphDef(context); case 'tf.pb.MetaGraphDef': return await openBinaryMetaGraphDef(context); case 'tf.pb.SavedModel': return await openSavedModel(context, await openBinarySavedModel(context), '', null); case 'tf.pb.FingerprintDef': return await openFingerprint(context); case 'tf.pb.mmap': return await openMemmapped(context); default: throw new tf.Error(`Unsupported TensorFlow format '${context.type}'.`); } } }; tf.Model = class { constructor(metadata, model, format, producer, bundle) { this.format = format; this.producer = producer || ''; this.modules = []; if (model) { for (let i = 0; i < model.meta_graphs.length; i++) { const meta_graph = model.meta_graphs[i]; let name = ''; if (meta_graph.meta_info_def && meta_graph.meta_info_def.any_info) { name = meta_graph.meta_info_def.any_info.toString(); } else if (model.meta_graphs.length > 1) { name = i.toString(); } const graph = new tf.Graph(metadata, meta_graph, name, bundle); this.modules.push(graph); } } else { const graph = new tf.Graph(metadata, null, '', bundle); this.modules.push(graph); } } }; tf.Graph = class { constructor(metadata, meta_graph, name, bundle) { this.name = name; this.nodes = []; this.inputs = []; this.outputs = []; this.functions = []; this.signatures = []; this.version = null; this.metadata = []; this.groups = false; if (meta_graph && meta_graph.graph_def) { const graph = meta_graph.graph_def; if (graph.versions) { this.version = `v${graph.versions.producer}`; } else if (graph.version) { this.version = graph.version; } else if (meta_graph.meta_info_def && meta_graph.meta_info_def.tensorflow_version) { this.version = meta_graph.meta_info_def.tensorflow_version; } if (meta_graph.meta_info_def && Array.isArray(meta_graph.meta_info_def.tags) && meta_graph.meta_info_def.tags.length > 0) { this.metadata.push(new tf.Argument('tags', meta_graph.meta_info_def.tags.join(', '))); } const output_arg_map = new Map(); metadata = new tf.GraphMetadata(metadata, graph.library); this.functions = metadata.functions; const context = new tf.Context(); const resolveTensorInfoName = (tensor) => { if (tensor) { if (tensor.name) { return tensor.name; } if (tensor.coo_sparse && tensor.coo_sparse.values_tensor_name) { return tensor.coo_sparse.values_tensor_name; } if (tensor.composite_tensor && Array.isArray(tensor.composite_tensor.components) && tensor.composite_tensor.components.length > 0) { return resolveTensorInfoName(tensor.composite_tensor.components[0]); } } return ''; }; for (const [key, signature_def] of Object.entries(meta_graph.signature_def)) { const inputs = []; for (const [key, tensor] of Object.entries(signature_def.inputs)) { const type = new tf.TensorType(tensor.dtype, tensor.tensor_shape); const name = resolveTensorInfoName(tensor).replace(/:0$/, ''); const value = context.value(name, type); const argument = new tf.Argument(key, [value]); inputs.push(argument); } const outputs = []; for (const [key, tensor] of Object.entries(signature_def.outputs)) { const type = new tf.TensorType(tensor.dtype, tensor.tensor_shape); const name = resolveTensorInfoName(tensor).replace(/:0$/, ''); const value = context.value(name, type); const argument = new tf.Argument(key, [value]); outputs.push(argument); output_arg_map.set(name, key); } const signature = new tf.Signature(key, inputs, outputs); this.signatures.push(signature); } const nodes = graph.node || []; context.graph(metadata, nodes, output_arg_map); this.nodes = context.nodes; this.inputs = context.inputs; this.outputs = context.outputs; } else if (bundle) { const nodes = new Map(); for (const tensor of bundle.tensors) { const parts = tensor.name.split('/'); if (bundle.format === 2) { if (tensor.name === '_CHECKPOINTABLE_OBJECT_GRAPH' || tensor.name.startsWith('optimizer/') || tensor.name.startsWith('keras_api/metrics/') || tensor.name.endsWith('/ExponentialMovingAverage') || tensor.name.indexOf('.OPTIMIZER_SLOT') !== -1) { continue; } if (tensor.name.endsWith('/.ATTRIBUTES/VARIABLE_VALUE')) { parts.pop(); parts.pop(); } } const tensorName = parts.pop(); const name = parts.join('/'); if (!nodes.has(name)) { nodes.set(name, []); } nodes.get(name).push({ name: tensorName, value: tensor }); } const namespaces = new Set(); this.nodes = Array.from(nodes).map(([name, value]) => { const node = { op: 'Node', name }; return new tf.Node(metadata, node, namespaces, new tf.Context(), value); }); } } }; tf.Signature = class { constructor(name, inputs, outputs) { this.name = name; this.inputs = inputs; this.outputs = outputs; } }; tf.Argument = class { constructor(name, value, type = null, visible = true) { this.name = name; this.value = value; this.type = type; this.visible = visible; } }; tf.Value = class { constructor(name, type, initializer = null) { if (typeof name !== 'string') { throw new tf.Error(`Invalid value identifier '${JSON.stringify(name)}'.`); } this.name = name; this.type = !type && initializer ? initializer.type : type; this.initializer = initializer; } }; tf.Function = class { constructor(metadata, name, func) { this.type = 'function'; this.name = name; this.version = null; this.tags = null; this.nodes = []; this.inputs = []; this.outputs = []; this.description = func ? null : 'Function definition not found.'; this.groups = false; const context = new tf.Context(); const input_arg = func && func.signature ? func.signature.input_arg : []; const output_arg = func && func.signature ? func.signature.output_arg : []; const ret = func && func.ret ? func.ret : {}; const nodes = func && func.node_def ? func.node_def : []; if (input_arg) { for (const input of input_arg) { const value = context.value(input.name, new tf.TensorType(input.type, null), null); const argument = new tf.Argument(input.name, [value]); this.inputs.push(argument); } } const output_arg_map = new Map(); if (output_arg) { const ret_map = new Map(); for (const key of Object.keys(ret)) { const value = func.ret[key]; const split = value.split(':', 2); ret_map.set(key, split[0]); } for (const output of output_arg) { const name = ret_map.get(output.name); const type = new tf.TensorType(output.type, null); const value = context.value(name, type, null); const argument = new tf.Argument(output.name, [value]); this.outputs.push(argument); output_arg_map.set(name, output.name); } } context.graph(metadata, nodes, output_arg_map); this.nodes = context.nodes; } }; tf.Node = class { constructor(metadata, node, namespaces, context, tensors) { this.type = node.metadata || metadata.type(node.op) || { name: node.op }; this.name = node.name; this.attributes = []; this.inputs = []; this.outputs = []; this.group = ''; if (node.name) { if (namespaces.has(node.name)) { this.group = node.name; } else { const index = node.name.lastIndexOf('/'); if (index !== -1) { const namespace = node.name.substring(0, index); if (namespaces.has(namespace)) { this.group = namespace; } } } } if (tensors) { for (const tensor of tensors) { const value = context.value(tensor.value.name, null, tensor.value); const argument = new tf.Argument(tensor.name, [value]); this.inputs.push(argument); } } else { if (node.device !== undefined) { this.device = node.device; } if (node.attr) { this.attributes = Object.entries(node.attr).map(([name, obj]) => { const schema = obj && obj.metadata ? obj.metadata : metadata.attribute(node.op, name); let value = null; let type = schema && typeof schema.type === 'string' ? schema.type : null; let visible = metadata.visible(node.op, name); switch (obj.value) { case undefined: type = ''; value = null; break; case 'type': type = 'type'; value = tf.Utility.dataType(obj.type); break; case 'i': value = obj.i; break; case 'f': value = obj.f; break; case 'b': value = obj.b; break; case 'shape': type = 'shape'; value = new tf.TensorShape(obj.shape); break; case 's': value = tf.Utility.decodeText(obj.s); break; case 'tensor': { type = 'tensor'; value = new tf.Tensor(obj.tensor); break; } case 'func': { type = 'function'; value = metadata.type(obj.func.name); // type = 'object'; // value = new tf.Node(metadata, { op: obj.func.name, attr: obj.func.attr }, null, new tf.Context()); break; } case 'placeholder': { type = 'placeholder'; value = obj; break; } case 'list': { const list = obj.list; if (list.s && list.s.length > 0) { value = list.s.map((s) => tf.Utility.decodeText(s)); } else if (list.i && list.i.length > 0) { value = list.i; } else if (list.f && list.f.length > 0) { value = list.f; } else if (list.type && list.type.length > 0) { type = 'type[]'; value = list.type.map((type) => tf.Utility.dataType(type)); } else if (list.shape && list.shape.length > 0) { type = 'shape[]'; value = list.shape.map((shape) => new tf.TensorShape(shape)); } else if (list.func && list.func.length > 0) { type = 'function[]'; value = list.func.map((func) => new tf.Node(metadata, { op: func.name, attr: func.attr })); } else { value = []; } break; } default: { throw new tf.Error(`Unsupported attribute value type '${JSON.stringify(value).substring(0, 32)}'.`); } } if (schema) { if (schema.visible === false) { visible = false; } else if (schema.default !== undefined) { const equals = (value, defaultValue) => { if (!Array.isArray(defaultValue) && defaultValue === Object(defaultValue)) { switch (defaultValue.type) { case 'type': defaultValue = tf.Utility.dataType(defaultValue.value); break; case 'shape': case 'tensor': defaultValue = defaultValue.value; break; default: throw new tf.Error(JSON.stringify(defaultValue)); } } if (typeof value === 'boolean' || typeof value === 'number' || typeof value === 'string') { return value === defaultValue; } if (typeof value === 'bigint') { return Number(value) === defaultValue; } return false; }; const defaultValue = schema.default; if (Array.isArray(value) && Array.isArray(defaultValue)) { if (value.length === defaultValue.length && value.every((item, index) => equals(item, defaultValue[index]))) { visible = false; } } else if (equals(value, defaultValue)) { visible = false; } } } if (name === '_class' || name === '_output_shapes' || visible === false) { visible = false; } return new tf.Argument(name, value, type, visible); }); } let inputIndex = 0; const inputs = (node.input || []).filter((input) => !input.name.startsWith('^')); if (this.type && this.type.inputs) { for (const input of this.type.inputs) { let count = 1; if (input.numberAttr) { const inputNumber = node.attr[input.numberAttr]; if (inputNumber && inputNumber.i) { count = Number(inputNumber.i); } } else if (input.typeListAttr) { const inputTypeListAttr = node.attr[input.typeListAttr]; if (inputTypeListAttr && inputTypeListAttr.list && inputTypeListAttr.list.type) { count = inputTypeListAttr.list.type.length; } } const values = inputs.slice(inputIndex, inputIndex + count).map((input) => context.value(input.name, null, null)); const argument = new tf.Argument(input.name, values); this.inputs.push(argument); inputIndex += count; } } this.inputs.push(...inputs.slice(inputIndex).map((input, index) => { const name = input.label ? input.label : (inputIndex + index).toString(); return new tf.Argument(name, [context.value(input.name)]); })); let outputIndex = 0; const outputs = node.output || []; if (this.type && this.type.outputs) { for (const output of this.type.outputs) { let count = 1; if (output.numberAttr) { const outputNumber = node.attr[output.numberAttr]; if (outputNumber && outputNumber.i) { count = Number(outputNumber.i); } } else if (output.typeListAttr) { const outputTypeListAttr = node.attr[output.typeListAttr]; if (outputTypeListAttr && outputTypeListAttr.list && outputTypeListAttr.list.type) { count = outputTypeListAttr.list.type.length; } } const values = outputs.slice(outputIndex, outputIndex + count).map((output) => { return context.value(output.name ? output.name : '-', null, null); }); const name = output.name ? output.name : `output${this.outputs.length === 0 ? '' : this.outputs.length}`; const argument = new tf.Argument(name, values); this.outputs.push(argument); outputIndex += count; } } this.outputs.push(...outputs.slice(outputIndex).map((output, index) => { const name = (outputIndex + index).toString(); const value = context.value(output.name ? output.name : '-', null, null); return new tf.Argument(name, [value]); })); const controlDependencies = node.controlDependencies || []; this.controlDependencies = controlDependencies.map((input) => context.value(input.name)); } } }; tf.Tensor = class { constructor(tensor, name, category = null) { this.name = name; this.category = category; if (tensor) { this.type = new tf.TensorType(tensor.dtype, tensor.tensor_shape || tensor.tensorShape); this._tensor = tensor; if (Object.prototype.hasOwnProperty.call(tensor, 'tensor_content')) { this._values = tensor.tensor_content; this.encoding = '<'; } else { const DataType = tf.proto.tensorflow.DataType; switch (tensor.dtype) { case DataType.DT_INVALID: { break; } case DataType.DT_BFLOAT16: { const values = tensor.half_val || []; this._values = new Uint8Array(values.length << 2); const view = new DataView(this._values.buffer, this._values.byteOffset, this._values.byteLength); for (let i = 0; i < values.length; i++) { view.setUint32(i << 2, values[i] << 16, true); } this.encoding = '<'; break; } case DataType.DT_HALF: { const values = tensor.half_val || []; this._values = new Uint8Array(values.length << 1); const view = new DataView(this._values.buffer, this._values.byteOffset, this._values.byteLength); for (let i = 0; i < values.length; i++) { view.setUint16(i << 1, values[i], true); } this.encoding = '<'; break; } case DataType.DT_FLOAT: { this._values = tensor.float_val || null; this.encoding = '|'; break; } case DataType.DT_DOUBLE: { this._values = tensor.double_val || null; this.encoding = '|'; break; } case DataType.DT_UINT8: case DataType.DT_UINT16: case DataType.DT_INT8: case DataType.DT_INT16: case DataType.DT_INT32: { this._values = tensor.int_val || null; this.encoding = '|'; break; } case DataType.DT_UINT32: { this._values = tensor.uint32_val || null; this.encoding = '|'; break; } case DataType.DT_INT64: { this._values = tensor.int64_val || null; this.encoding = '|'; break; } case DataType.DT_UINT64: { this._values = tensor.uint64_val || null; this.encoding = '|'; break; } case DataType.DT_BOOL: { this._values = tensor.bool_val || null; this.encoding = '|'; break; } case DataType.DT_STRING: { this._values = tensor.string_val || null; this.encoding = '|'; break; } case DataType.DT_COMPLEX64: { const values = tensor.scomplex_val || null; this._values = new Array(values.length >> 1); for (let i = 0; i < values.length; i += 2) { this._values[i >> 1] = new base.Complex(values[i], values[i + 1]); } this.encoding = '|'; break; } case DataType.DT_COMPLEX128: { const values = tensor.dcomplex_val || null; this._values = new Array(values.length >> 1); for (let i = 0; i < values.length; i += 2) { this._values[i >> 1] = new base.Complex(values[i], values[i + 1]); } this.encoding = '|'; break; } case DataType.DT_FLOAT8_E5M2: case DataType.DT_FLOAT8_E4M3FN: case DataType.DT_FLOAT8_E4M3FNUZ: case DataType.DT_FLOAT8_E4M3B11FNUZ: case DataType.DT_FLOAT8_E5M2FNUZ: { this._values = tensor.float8_val || null; this.encoding = '<'; break; } default: { throw new tf.Error(`Unsupported tensor data type '${tensor.dtype}'.`); } } } } else { this.type = new tf.TensorType('?', null); this._tensor = null; } } get values() { let values = this._values; if (this.encoding === '|' && Array.isArray(values)) { if (this.type.dataType === 'string') { values = values.map((value) => tf.Utility.decodeText(value)); } const shape = (this._tensor.tensor_shape || this._tensor.tensorShape).dim.map((dim) => dim.size); const size = shape.reduce((a, b) => a * Number(b), 1); if (values.length === 1 && size > 1) { values = new Array(size).fill(values[0]); } } return values; } }; tf.TensorType = class { constructor(dtype, shape) { this.dataType = dtype ? tf.Utility.dataType(dtype) : '?'; this.shape = new tf.TensorShape(shape); } equals(obj) { return obj && this.dataType === obj.dataType && this.shape.equals(obj.shape); } toString() { return this.dataType + this.shape.toString(); } }; tf.TensorShape = class { constructor(shape) { this.dimensions = null; if (shape) { if (shape.unknown_rank) { this.dimensions = null; } else if (Array.isArray(shape.dim)) { if (shape.dim.length === 0) { this.dimensions = []; } else if (shape.dim.length === 1 && !shape.dim[0].size) { this.dimensions = [0]; } else { this.dimensions = shape.dim.map((dim) => { const size = dim.size && dim.size.toNumber ? dim.size.toNumber() : dim.size; return size && size !== -1 ? size : '?'; }); } } } } equals(obj) { return (this.dimensions === null && obj.dimensions === null) || (Array.isArray(this.dimensions) && Array.isArray(obj.dimensions) && this.dimensions.length === obj.dimensions.length && this.dimensions.every((value, index) => obj.dimensions[index] === value)); } toString() { if (this.dimensions === null) { return '[?]'; } if (this.dimensions.length === 0) { return ''; } return `[${this.dimensions.map((dim) => (dim && dim !== -1) ? dim.toString() : '?').join(',')}]`; } }; tf.TensorBundle = class { static async open(stream, identifier, context) { const format = identifier.toLowerCase().endsWith('.index') ? 2 : 1; const table = new tf.TensorBundle.Table(stream); if (!table.entries.has('')) { throw new tf.Error('Bundle header not available.'); } if (format === 1) { return new tf.TensorBundle(format, table.entries, []); } const buffer = table.entries.get(''); const reader = protobuf.BinaryReader.open(buffer); const header = tf.proto.tensorflow.BundleHeaderProto.decode(reader); const numShards = header.num_shards; const promises = []; for (let i = 0; i < numShards; i++) { const shardIndex = (`0000${i}`).slice(-5); const shardCount = (`0000${numShards}`).slice(-5); const filename = identifier.split('.'); filename.pop(); const basename = filename.join('.'); const name = `${basename}.data-${shardIndex}-of-${shardCount}`; promises.push(context.fetch(name)); } try { const contexts = await Promise.all(promises); const streams = contexts.map((context) => context.stream); return new tf.TensorBundle(format, table.entries, streams); } catch (error) { context.error(error, false); return new tf.TensorBundle(format, table.entries, null); } } constructor(format, entries, streams) { this.format = format; this.tensors = []; switch (format) { case 1: { const buffer = entries.get(''); const reader = protobuf.BinaryReader.open(buffer); const header = tf.proto.tensorflow.SavedTensorSlices.decode(reader); const data = new Map(); for (const [name, buffer] of entries) { if (name !== '' && name !== 'global_step') { const reader = protobuf.BinaryReader.open(buffer); const slices = tf.proto.tensorflow.SavedTensorSlices.decode(reader); const name = slices.data.name; const tensor = slices.data.data; if (data.has(name)) { const item = data.get(name); if (item !== null) { if (tensor[item.key] && tensor[item.key].length > 0) { item.value = item.value.concat(tensor[item.key]); } else { data.set(name, null); } } } else if (tensor.tensor_content && tensor.tensor_content.length > 0) { data.set(name, { key: 'tensor_content', value: tensor.tensor_content }); } else { const keys = Object.keys(tensor).filter((key) => key.endsWith('_val') && tensor[key] && tensor[key].length > 0); data.set(name, keys.length === 1 ? { key: keys[0], value: tensor[keys[0]] } : null); } } } for (const meta of header.meta.tensor) { if (meta.name !== 'global_step') { const tensor = new tf.proto.tensorflow.TensorProto(); tensor.dtype = meta.type; tensor.tensor_shape = meta.shape; const item = data.get(meta.name); if (item) { tensor[item.key] = item.value; } this.tensors.push(new tf.Tensor(tensor, meta.name, null)); } } break; } case 2: { entries.forEach((buffer, name) => { if (name !== '') { const reader = protobuf.BinaryReader.open(buffer); const entry = tf.proto.tensorflow.BundleEntryProto.decode(reader); const tensor = new tf.proto.tensorflow.TensorProto(); tensor.dtype = entry.dtype; tensor.tensor_shape = entry.shape; const offset = typeof entry.offset === 'bigint' ? Number(entry.offset) : entry.offset; const size = typeof entry.size === 'bigint' ? Number(entry.size) : entry.size; if (streams) { const stream = streams[entry.shard_id]; stream.seek(offset); tensor.tensor_content = stream.peek(size); } this.tensors.push(new tf.Tensor(tensor, name, null)); } }); break; } default: { throw new tf.Error(`Unsupported Tensor Bundle format '${format}'.`); } } } }; tf.TensorBundle.Table = class { constructor(stream) { // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/lib/io/table.cc this.entries = new Map(); if (stream.length <= 54) { throw new tf.Error('Invalid index file size.'); } stream.seek(-48); const buffer = stream.peek(48); const reader = new tf.BinaryReader(buffer); reader.seek(-8); const signature = [0x57, 0xfb, 0x80, 0x8b, 0x24, 0x75, 0x47, 0xdb]; if (!reader.read(8).every((value, index) => value === signature[index])) { throw new tf.Error('Invalid table signature.'); } reader.seek(-48); // kEncodedLength reader.varint64(); // metaindex offset reader.varint64(); // metaindex size const indexOffset = reader.varint64(); const indexSize = reader.varint64(); const indexBlock = new tf.TensorBundle.Table.Block(stream, indexOffset, indexSize); for (const [, value] of indexBlock.entries) { const valueReader = new tf.BinaryReader(value); const offset = valueReader.varint64(); const size = valueReader.varint64(); const block = new tf.TensorBundle.Table.Block(stream, offset, size); for (const [name, value] of block.entries) { this.entries.set(name, value); } } stream.seek(0); } }; tf.TensorBundle.Table.Block = class { constructor(stream, offset, size) { // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/lib/io/block.cc this.entries = new Map(); stream.seek(offset); const buffer = stream.read(size); // blockContents const [compression] = stream.read(1); stream.skip(4); // crc32 let reader = new tf.BinaryReader(buffer); switch (compression) { case 0: // kNoCompression break; case 1: // kSnappyCompression reader = new tf.BinaryReader(reader.unsnappy()); break; default: throw new tf.Error(`Unsupported block compression '${compression}'.`); } reader.seek(-4); const numRestarts = reader.int32(); reader.seek(-4 - (4 * numRestarts)); const restartOffsets = []; for (let i = 0; i < numRestarts; i++) { restartOffsets.push(reader.int32()); } const decoder = new TextDecoder('utf-8'); for (let i = 0; i < numRestarts; i++) { reader.seek(restartOffsets[i]); let key = ''; while (reader.position < reader.length) { const sharedSize = reader.varint32(); // index shared size const nonSharedSize = reader.varint32(); // index non shared size const valueSize = reader.varint32(); if (sharedSize === 0 && nonSharedSize === 0 && valueSize === 0) { break; } key = key.substring(0, sharedSize); key += decoder.decode(reader.read(nonSharedSize)); const value = reader.read(valueSize); this.entries.set(key, value); } } } }; tf.BinaryReader = class { constructor(buffer) { this._reader = base.BinaryReader.open(buffer); this._decoder = new TextDecoder('utf-8'); } get length() { return this._reader.length; } get position() { return this._reader.position; } seek(position) { this._reader.seek(position); } read(length) { return this._reader.read(length); } byte() { return this._reader.byte(); } int32() { return this._reader.int32(); } uint32() { return this._reader.uint32(); } string() { const size = this.uint32(); const buffer = this.read(size); return this._decoder.decode(buffer); } varint32() { return this.varint64(); } varint64() { let result = 0; for (let shift = 0; shift <= 63; shift += 7) { const byte = this.byte(); if (byte & 128) { result |= (byte & 127) << shift; } else { result |= byte << shift; break; } } return result; } unsnappy() { const data = new Uint8Array(this.varint64()); const mask = [0, 0xff, 0xffff, 0xffffff, 0xffffffff]; let position = 0; while (this._position < this._length) { let length = 0; const c = this.byte(); switch (c & 0x03) { case 0: { length = (c >>> 2) + 1; if (length > 60) { const short = length - 60; length = (this.uint32() & mask[short]) + 1; this._position += short - 4; } data.set(this.read(length), position); break; } case 1: { length = ((c >>> 2) & 0x07) + 4; const offset = this.byte() + ((c >>> 5) << 8); data.set(data.subarray(position - offset, position - offset + length), position); break; } case 2: { length = (c >>> 2) + 1; const offset = this.uint16(); data.set(data.subarray(position - offset, position - offset + length), position); break; } case 3: { length = (c >>> 2) + 1; const offset = this.uint32(); data.set(data.subarray(position - offset, position - offset + length), position); break; } default: { break; } } position += length; } return data; } }; tf.EventFileReader = class { static open(stream) { if (stream.length < 16) { return null; } const masked_crc32c = (bytes) => { const poly = 0x82f63b78; let crc = 0xffffffff; for (let n = 0; n < bytes.length; n++) { crc ^= bytes[n]; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc = crc & 1 ? (crc >>> 1) ^ poly : crc >>> 1; crc >>>= 0; } crc ^= 0xffffffff; crc >>>= 0; crc = ((crc >> 15) | (crc << 17)) + 0xa282ead8; crc >>>= 0; return crc; }; const buffer = stream.peek(12); const reader = new tf.BinaryReader(buffer); const length_bytes = reader.read(8); const length_crc = reader.uint32(); if (masked_crc32c(length_bytes) !== length_crc) { return null; } return new tf.EventFileReader(stream); } constructor(stream) { this._stream = stream; } read() { if (this._stream.position < this._stream.length) { const uint64 = (stream) => { const buffer = stream.read(8); const view = new DataView(buffer.buffer, buffer.byteOffset, buffer.byteLength); const value = view.getBigUint64(0, true); return value.toNumber(); }; const length = uint64(this._stream); this._stream.skip(4); // masked crc of length const buffer = this._stream.read(length); const reader = protobuf.BinaryReader.open(buffer); const event = tf.proto.tensorflow.Event.decode(reader); this._stream.skip(4); // masked crc of data return event; } return null; } }; tf.GraphMetadata = class { constructor(metadata, library) { this._metadata = metadata; this._functions = new Map(); this._attributes = new Map(); this._visibleCache = new Map(); if (library && Array.isArray(library.function) && library.function.length > 0) { for (const func of library.function) { const name = func.signature.name; if (this._functions.has(func.name)) { throw new tf.Error(`Duplicate function name '${func.name}'.`); } this._functions.set(name, func); } } } type(name) { if (this._functions.has(name)) { const func = this._functions.get(name); if (func instanceof tf.Function) { return func; } this._functions.set(name, new tf.Function(this, func.signature.name, func)); return this._functions.get(name); } const type = this._metadata.type(name); if (!type) { this._functions.set(name, new tf.Function(this, name, null)); return this._functions.get(name); } return type; } attribute(type, name) { const key = `${type}::${name}`; if (!this._attributes.has(key)) { const schema = this.type(type); if (schema && schema.attributes) { for (const attribute of schema.attributes) { const key = `${type}::${attribute.name}`; this._attributes.set(key, attribute); } } } return this._attributes.get(key); } visible(type, name) { if (!this._visibleCache.has(type)) { const set = new Set(); const schema = this.type(type); if (schema && schema.inputs) { for (const input of schema.inputs) { if (input.typeAttr) { set.add(input.typeAttr); } else if (input.typeListAttr) { set.add(input.typeListAttr); } if (input.numberAttr) { set.add(input.numberAttr); } } } if (schema && schema.outputs) { for (const output of schema.outputs) { if (output.typeAttr) { set.add(output.typeAttr); } else if (output.typeListAttr) { set.add(output.typeListAttr); } if (output.numberAttr) { set.add(output.numberAttr); } } } this._visibleCache.set(type, set); } return !this._visibleCache.get(type).has(name); } get functions() { for (const [name, func] of this._functions) { if (func instanceof tf.Function === false) { this._functions.set(name, new tf.Function(this, func.signature.name, func)); } } return Array.from(this._functions.values()); } }; tf.Context = class { constructor() { this._values = new Map(); this.signatures = []; this.nodes = []; } value(name, type, tensor) { if (name.length === 0 && tensor) { return new tf.Value(name, type || null, tensor); } if (!this._values.has(name)) { this._values.set(name, new tf.Value(name, type || null, tensor || null)); } else if ((type && !type.equals(this._values.get(name).type)) || tensor) { throw new tf.Error(`Duplicate value '${name}'.`); } return this._values.get(name); } graph(metadata, nodes, output_arg_map) { const namespaces = new Set(); nodes = new Map(nodes.map((node) => [node.name, node])); this.inputs = []; this.outputs = []; for (const [name, node] of nodes) { if (node.op !== 'Const') { const index = name.lastIndexOf('/'); if (index !== -1) { const namespace = name.substring(0, index); namespaces.add(namespace); } } node.output = []; } const node_output = (input) => { let name = input; let index = 0; const control = name.startsWith('^'); if (control) { name = name.substring(1); } const colon = name.lastIndexOf(':'); if (colon !== -1) { const suffix = name.substring(colon + 1); const candidate = name.substring(0, colon); const value = parseInt(suffix, 10); if (!isNaN(value) && nodes.has(candidate) && !nodes.has(name)) { index = value; name = candidate; } } const from = nodes.get(name); if (from) { for (let i = from.output.length; i <= index; i++) { const key = i === 0 ? from.name : `${from.name}:${i}`; const value = { name: key, to: [] }; from.output.push(value); } } const key = index === 0 ? name : `${name}:${index}`; return [key, index, control, from]; }; for (const node of nodes.values()) { const inputs = node.input; node.input = []; node.controlDependencies = []; for (const input of inputs) { const [key, index, control, from] = node_output(input); if (from) { from.output[index].to.push(node); } const value = { name: key, from }; if (control) { node.controlDependencies.push(value); } else { node.input.push(value); } } } if (output_arg_map) { for (const [name, node] of nodes) { if (output_arg_map.has(name)) { node.output.push({ name, to: [] }); } } } const map_tensor = (name, node, kind) => { if (node && node.op === 'Const' && node.input.length === 0 && node.output.length === 1 && node.output[0].to.length === 1 && node.controlDependencies.length === 0) { const value = node.attr.value; if (value && Object.prototype.hasOwnProperty.call(value, 'tensor')) { const tensor = new tf.Tensor(value.tensor, name, kind); return this.value(name, tensor.type, tensor); } } return null; }; const map_resource = (name, node, tensor) => { if (node && node.op === 'Placeholder' && node.input.length === 0 && node.output.length === 1 && node.controlDependencies.length === 0) { const dtype = node.attr.dtype.type; if (dtype === tf.proto.tensorflow.DataType.DT_RESOURCE) { return this.value(name, null, tensor); } } return null; }; for (const node of nodes.values()) { if (node.op === 'Identity' && node.input.length === 1 && node.output.length === 1 && node.output[0].to.length === 1 && node.controlDependencies.length === 0) { const initializer = map_tensor(node.name, node.input[0].from, 'Identity Constant'); if (initializer) { nodes.delete(initializer.name); nodes.delete(node.input[0].name); } const identity = node.input[0].from; if (identity && identity.op === 'Identity' && identity.input.length === 1 && identity.output.length === 1 && node.output[0].to.length === 1 && node.controlDependencies.length === 0) { const initializer = map_tensor(node.name, identity.input[0].from, 'Identity Constant'); if (initializer) { nodes.delete(initializer.name); nodes.delete(initializer.name); nodes.delete(identity.name); nodes.delete(node.name); } } } } for (const node of nodes.values()) { const initializer = map_tensor(node.name, node, 'Const'); if (initializer) { nodes.delete(node.name); nodes.delete(initializer.name); } } for (const node of nodes.values()) { if (node.op === 'ReadVariableOp' && node.input.length === 1 && node.output.length === 1 && node.output[0].to.length === 1 && node.controlDependencies.length === 0) { if (node.attr && node.attr.dtype && node.attr._output_shapes && node.attr._output_shapes.list && node.attr._output_shapes.list.shape) { const tensor = new tf.proto.tensorflow.TensorProto(); tensor.dtype = node.attr.dtype.type; [tensor.tensor_shape] = node.attr._output_shapes.list.shape; const name = node.name; const initializer = map_resource(name, node.input[0].from, new tf.Tensor(tensor, name, 'Resource Variable')); if (initializer) { nodes.delete(initializer.name); nodes.delete(node.input[0].name); } } } } const inputs = new Map(); for (const [name, node] of nodes) { if (node.op === 'Placeholder' && node.attr && node.attr.dtype && Number.isInteger(node.attr.dtype.type) && node.attr._output_shapes && node.attr._output_shapes.list && Array.isArray(node.attr._output_shapes.list.shape) && node.attr._output_shapes.list.shape.length > 0 && node.input.length === 0 && node.output.length === 1 && node.controlDependencies.length === 0) { const type = new tf.TensorType(node.attr.dtype.type, node.attr._output_shapes.list.shape[0]); const value = this.value(name, type, null); const argument = new tf.Argument(name, [value]); inputs.set(name, argument); nodes.delete(name); } } const updateTorchScript = (nodes) => { for (const node of nodes.values()) { if (node.op === 'prim::Constant' && node.input.length === 0 && node.controlDependencies.length === 0 && node.attr && Object.keys(node.attr).length === 1 && node.attr.attr && node.attr.attr.s) { const value = tf.Utility.decodeText(node.attr.attr.s); const match = /{\s*value\s*:\s*(.*)\s*}/.exec(value); if (match) { node.value = match[1].trim(); } const empty = /{\s*}/.exec(value); if (empty) { node.value = null; } } if (node.op === 'prim::GetAttr' && node.input.length === 1 && node.controlDependencies.length === 0 && node.attr && Object.keys(node.attr).length === 1 && node.attr.attr && node.attr.attr.s) { const value = tf.Utility.decodeText(node.attr.attr.s); const match = /{\s*name\s*:\s*([A-Za-z0-9_]*)\s*}/.exec(value); if (match) { node.value = match[1].trim(); } } if (node.op === 'IO Node' && node.controlDependencies.length === 0) { const shape = node.attr && node.attr._output_shapes && node.attr._output_shapes.list && node.attr._output_shapes.list.shape ? node.attr._output_shapes.list.shape[0] : null; const type = shape ? new tf.TensorType('?', shape) : null; if (node.input.length === 0 && node.output.length === 1) { const argument = new tf.Argument(node.name, [this.value(node.output[0].name, type, null)]); this.inputs.push(argument); nodes.delete(node.name); } if (node.input.length === 1 && node.output.length === 0) { const argument = new tf.Argument(node.name, [this.value(node.input[0].name, type, null)]); this.outputs.push(argument); nodes.delete(node.name); } } if (Object.keys(node.attr).length === 2 && node.attr.attr && node.attr.attr.s && node.attr._output_shapes) { const value = tf.Utility.decodeText(node.attr.attr.s); if (/\s*/.exec(value) || /{\s*}/.exec(value)) { node.attr = {}; delete node._output_shapes; } } } const remove_input = (input, node) => { const from = input.from; if (from) { for (const output of from.output) { output.to = output.to.filter((to) => to !== node); } if (from.output.every((output) => output.to.length === 0) && from.controlDependencies.length === 0) { from.remove = true; } delete input.from; } }; for (const node of nodes.values()) { if (node.op === 'prim::ListConstruct' && node.input.every((input) => input.from.value !== undefined) && node.controlDependencies.length === 0) { node.value = node.input.map((input) => input.from.value); for (const input of node.input) { remove_input(input, node); } node.input = []; } } for (const node of nodes.values()) { const remove = new Set(); for (let i = 0; i < node.input.length; i++) { const input = node.input[i]; const from = input.from; if (from) { if (from.op === 'prim::GetAttr' && from.input.length === 1 && from.output.length === 1 && from.controlDependencies.length === 0 && from.value !== undefined) { remove_input(input, node); input.label = from.value; const tensor = new tf.Tensor(null, input.name, from.op); this.value(input.name, null, tensor); } if (from.op === 'prim::Constant' && from.input.length === 0 && from.controlDependencies.length === 0 && from.value !== undefined) { input.constant = from.value; remove_input(input, node); remove.add(input.name); } if (from.op === 'prim::ListConstruct' && from.output.length === 1 && from.controlDependencies.length === 0 && from.value !== undefined) { input.list = from.value; remove_input(input, node); remove.add(input.name); } } } if (node.__metadata__) { const torch = node.__torch__; const match = (node, schema) => { const args = schema.arguments || []; const inputs = node.input || []; if (inputs.length > args.length) { return false; } for (let i = 0; i < inputs.length; i++) { const input = inputs[i]; const arg = args[i]; let type = arg.real_type; type = type instanceof torch.OptionalType ? type.getElementType() : type; switch (type.str()) { case 'Tensor': { if ((input.constant === undefined && input.list === undefined) || input.constant === null) { continue; } break; } case 'int': case 'SymInt': { if (input.constant !== undefined && Number.isInteger(parseInt(input.constant, 10))) { continue; } break; } case 'float': { if (input.constant !== undefined && !isNaN(parseFloat(input.constant))) { continue; } break; } case 'int[]': case 'int[2]': case 'SymInt[]': case 'SymInt[2]': { if (Array.isArray(input.list)) { const list = input.list.map((item) => parseInt(item, 10)); if (list.every((value) => Number.isInteger(value))) { continue; } } break; } case 'bool': { if (input.constant === 'false' || input.constant === 'true' || input.constant === '0' || input.constant === '1') { continue; } break; } case 'Scalar': { if (input.constant !== undefined && Number.isInteger(parseInt(input.constant, 10))) { continue; } break; } default: { break; } } return false; } return true; }; const schema = node.__metadata__.find((schema) => match(node, schema)); if (schema) { const args = schema.arguments; const inputs = node.input || []; for (let i = 0; i < inputs.length; i++) { const input = inputs[i]; delete input.metadata; const arg = args[i]; let type = arg.real_type; type = type instanceof torch.OptionalType ? type.getElementType() : type; switch (type.str()) { case 'Tensor': { input.metadata = arg; break; } case 'int': case 'SymInt': { const value = parseInt(input.constant, 10); input.attr = new tf.proto.tensorflow.AttrValue(); input.attr.i = value; input.attr.metadata = arg; break; } case 'float': { const value = parseFloat(input.constant, 10); input.attr = new tf.proto.tensorflow.AttrValue(); input.attr.f = value; input.attr.metadata = arg; break; } case 'int[]': case 'int[2]': case 'SymInt[]': case 'SymInt[2]': { const list = input.list.map((item) => parseInt(item, 10)); input.attr = new tf.proto.tensorflow.AttrValue(); input.attr.list = new tf.proto.tensorflow.ListValue(); input.attr.list.i = list; input.attr.metadata = arg; break; } case 'bool': { input.attr = new tf.proto.tensorflow.AttrValue(); input.attr.b = input.constant === 'true' || input.constant === '1'; input.attr.metadata = arg; break; } case 'Scalar': { const value = parseInt(input.constant, 10); input.attr = new tf.proto.tensorflow.AttrValue(); input.attr.i = value; input.attr.metadata = arg; break; } default: { break; } } } node.metadata = { ...schema }; node.metadata.name = node.op; } } node.input = node.input.filter((input, index) => { if (input.attr) { const name = input.attr.metadata ? input.attr.metadata.name : index.toString(); node.attr[name] = input.attr; } else if (input.constant !== undefined && input.constant !== null) { const attr = new tf.proto.tensorflow.AttrValue(); attr.s = input.constant; node.attr[index.toString()] = attr; } else if (input.list !== undefined) { const attr = new tf.proto.tensorflow.AttrValue(); attr.list = new tf.proto.tensorflow.ListValue(); attr.list.s = input.list; node.attr[index.toString()] = attr; } return !remove.has(input.name); }); } for (const node of nodes.values()) { if (node.op === 'prim::GetAttr' && node.remove) { nodes.delete(node.name); } if (node.op === 'prim::Constant' && node.remove) { nodes.delete(node.name); } if (node.op === 'prim::ListConstruct' && node.remove) { nodes.delete(node.name); } } }; updateTorchScript(nodes); for (const input of inputs.values()) { this.inputs.push(input); } for (const node of nodes.values()) { this.nodes.push(new tf.Node(metadata, node, namespaces, this)); } } }; tf.Utility = class { static decodeText(value) { if (typeof value === 'string') { return value; } if (value.length === 0) { return ''; } tf.Utility._utf8Decoder = tf.Utility._utf8Decoder || new TextDecoder('utf-8'); return tf.Utility._utf8Decoder.decode(value); } static dataType(type) { if (!tf.Utility._dataTypes) { const DataType = tf.proto.tensorflow.DataType; const dataTypes = new Map(Object.entries(DataType).map(([name, value]) => { const key = name.startsWith('DT_') ? name.substring(3) : name; return [value, key.toLowerCase()]; })); dataTypes.set(DataType.DT_HALF, 'float16'); dataTypes.set(DataType.DT_FLOAT, 'float32'); dataTypes.set(DataType.DT_DOUBLE, 'float64'); dataTypes.set(DataType.DT_BOOL, 'boolean'); dataTypes.set(DataType.DT_COMPLEX64, 'complex'); dataTypes.set(DataType.DT_COMPLEX128, 'complex'); dataTypes.set(DataType.DT_FLOAT8_E5M2, 'float8e5m2'); dataTypes.set(DataType.DT_FLOAT8_E4M3FN, 'float8e4m3fn'); dataTypes.set(DataType.DT_FLOAT8_E4M3FNUZ, 'float8e4m3fnuz'); dataTypes.set(DataType.DT_FLOAT8_E4M3B11FNUZ, 'float8e4m3b11fnuz'); dataTypes.set(DataType.DT_FLOAT8_E5M2FNUZ, 'float8e5m2fnuz'); tf.Utility._dataTypes = dataTypes; } return tf.Utility._dataTypes.has(type) ? tf.Utility._dataTypes.get(type) : '?'; } static dataTypeKey(type) { if (!tf.Utility._dataTypeKeys) { tf.Utility.dataType(0); tf.Utility._dataTypeKeys = new Map(Array.from(tf.Utility._dataTypes).map(([key, value]) => [value, key])); } return tf.Utility._dataTypeKeys.get(type); } }; tf.Error = class extends Error { constructor(message) { super(message); this.name = 'Error loading TensorFlow model.'; } }; export const ModelFactory = tf.ModelFactory;