Theory TensorFlow_Import
section‹Infrastructure for Importing TensorFlow Models›
theory
TensorFlow_Import
imports
Complex_Main
Nano_JSON.Nano_JSON_Main
keywords
"import_TensorFlow" :: thy_decl
and "as"::quasi_command
begin
text‹
In this theory, we implement the core infrastructure for importing models from
TensorFlow.js~\<^cite>‹"smilkov.ea:tensorflowjs:2019"›. This common infrastructure provided a generic
parser for the JSON~\<^cite>‹"ecma:json:2017" and "ietf:rfc8259-json:2017"› representation of
neural networks that can be exported from TensorFlow.js. Actually, TensorFlow.js~\<^cite>‹"smilkov.ea:tensorflowjs:2019"› exports the structure of a neural network (and its configuration
used for training the neural network) as JSON file. The weights and biases are stored
in a binary file to which the JSON file refers to (see 🌐‹https://www.tensorflow.org/js/guide/save_load›
and 🌐‹https://github.com/tensorflow/tfjs/issues/386› for more details).
This theory implements an infrastructure for importing this format, including the decoding
of the binary format storing the weights and biases, into Isabelle/HOL. At its core, the
infrastructure provides a parser for the format used by TensorFlow.js and a mechanism
for hooking datatype packages into it that provide specific encodings into Isabelle/HOL.
As a first example, this theory provides a datatype package that provides a JSON-like
encoding of neural networks (including their weights and biases) using Nano
JSON~\<^cite>‹"brucker:nano-json:2022"›. The implementation used the JSON infrastructure
provided by the AFP entry Nano JSON~\<^cite>‹"brucker:nano-json:2022"›.
›
subsection‹Encoder›
ML‹
signature TENSORFLOW_TYPE = sig
datatype activationT = Linear | BinaryStep | Softsign | Sign | Sigmoid | Swish | Tanh | Relu | Gelu
| GRelu | Softplus | Elu | Selu | Exponential | Hard_sigmoid
| Softmax | Softmax_taylor | Sigmoid_taylor
datatype layerT = Dense | InputLayer | OutputLayer
type 'a layer = {
activation: activationT option,
bias: 'a list, units: int,
layer_type: layerT,
name: string,
weights: 'a list list
}
end
›
ML_file‹Tools/TensorFlow_Type.ML›
text‹
The ML structure @{ML_structure ‹TensorFlow_Type:TENSORFLOW_TYPE› } provides the core datatypes
required for the TensorFlow.js import:
▪ @{ML_type ‹TensorFlow_Type.activationT›}: this datatype enumerates the currently supported
activation functions (see~\autoref{tab:tensorflow-activation} for a mapping of their
names used by our formalization).
▪ @{ML_type ‹TensorFlow_Type.layerT›}: This datatype enumerates the currently supported layer types
of TensorFlow.
▪ @{ML_type ‹'a TensorFlow_Type.layer›}: This record captures the properties of a layer that are
extracted from the JSON provided by TensorFlow.js.
›
ML‹
signature TENSORFLOW_JSON = sig
val transform_json: Path.T -> Nano_Json_Type.json -> Nano_Json_Type.json
val def_nn_json: bstring -> typ -> typ -> Nano_Json_Type.json
-> local_theory -> local_theory
val convert_layers: Nano_Json_Type.json
-> IEEEReal.decimal_approx TensorFlow_Type.layer list
end
›
ML_file‹Tools/TensorFlow_Json.ML›
text‹
TensorFlow.js does export a neural network in a format consisting out of a JSON file and and a
binary file:
▪ the JSON file stores the overall structure of the neural network and the configuration used for
training the neural network. Notably, the JSON file does neither contain the biases nor the weights.
▪ a binary file containing the biases and weights.
›
text‹
The ML structure @{ML_structure ‹TensorFlow_Json:TENSORFLOW_JSON› } provides, foremost, a function
for parsing the JSON exported neural network in the format supported by TensorFlow. This
function, @{ML ‹TensorFlow_Json.transform_json›}, takes two arguments
▪ the directory (path) of the TensorFlow.js export, i.e., the directory in which both the
JSON file and the binary file containing the biases and weights are stored.
▪ the parsed JSON file (the actual JSON parsing is done using @{ML "Nano_Json_Parser.json_of_string"},
which is provided by Nano JSON~\<^cite>‹"brucker:nano-json:2022"›.
›
text‹
The function @{ML ‹TensorFlow_Json.transform_json›} parses the binary file containing the biases
and weights and transforms the input JSON such that the resulting JSON representation includes
the biases and weights. In more detail, the JSON file exported by TensorFlow.js stores the
biases and weights as follows:
\begin{json}
"weightsManifest": [{
"paths": [ "group1-shard1of1.bin" ],
"weights": [
{
"name": "dense/kernel",
"shape": [2, 1],
"dtype": "float32"
}, {
"name": "dense/bias",
"shape": [1],
"dtype": "float32"
}
]
}]
\end{json}
Instead of storing the biases and weights in the JSON file, the exported JSON only contains
the type information (here: \inlinejson|float32|) refers to an external file (here:
\inlinejson|group1-shard1of1.bin| that stores the actual value. In our example, this
external file has a size of 12 bytes, storing three 32 Bit floating point numbers (encoding
as IEEE floating point using a Little Endian encoding. The order of the biases and weights
corresponds to the order and shape of their references in the original JSON file. In our
example, the function @{ML ‹TensorFlow_Json.transform_json›} results in the following transformed
\inlinejson{weightsManifest}:
\begin{json}
"weightsManifest":[{
"name":"dense/kernel",
"shape":[
[-0.47318925857543945E1],
[-0.4610690593719482E1]
]
},
{
"name":"dense/bias",
"shape":[[0.22737088203430176E1]]
}]
\end{json}
Moreover, the ML structure @{ML_structure ‹TensorFlow_Json:TENSORFLOW_JSON› } also provides an
ML for converting a (transformed) JSON representation into a more abstract representation based on
a sequence of layers: @{ML ‹TensorFlow_Json.convert_layers›}. This function uses the datatypes
provided by @{ML_structure ‹TensorFlow_Type:TENSORFLOW_TYPE›}.
Finally, @{ML_structure ‹TensorFlow_Json:TENSORFLOW_JSON› } provides
@{ML ‹TensorFlow_Json.def_nn_json›}, which is a simple wrapper around the datatype package provided
by Nano JSON generating a formal JSON representation of the neural network imported from TensorFlow.js
in HOL.
›
ML‹
signature CONVERT_TENSORFLOW_SYMTAB = sig
structure NN_Encoder_Data: THEORY_DATA
eqtype targetT
type nn_encoderT = bstring-> typ -> typ -> Nano_Json_Type.json -> local_theory -> local_theory
val add_encoding: Symtab.key * nn_encoderT -> theory -> theory
val assert_target: theory -> Symtab.key -> Symtab.key
val lookup_nn_encoder: theory -> Symtab.key -> (targetT * nn_encoderT) option
end
›
ML_file‹Tools/Convert_TensorFlow_Symtab.ML›
text‹
We use the mechanism of attaching a symtab to theories to provide a dynamic registration
mechanism for different datatype packages that encode the JSON representation in a formal
model. The ML structure @{ML_structure ‹Convert_TensorFlow_Symtab:CONVERT_TENSORFLOW_SYMTAB› }
defines the type for encoder functions (i.e., @{ML_type ‹Convert_TensorFlow_Symtab.nn_encoderT›} and
it provides methods for adding a new encoding (@{ML ‹Convert_TensorFlow_Symtab.add_encoding›},
checking if a encoder for a given target encoding exists (@{ML ‹Convert_TensorFlow_Symtab.assert_target›},
and for the lookup of an encoder ((@{ML ‹Convert_TensorFlow_Symtab.lookup_nn_encoder›}). The
symtab is registered as follows:
›
ML‹val _ = Theory.setup
(Convert_TensorFlow_Symtab.add_encoding
("json", TensorFlow_Json.def_nn_json))
›
ML‹
local
fun import_and_define_nn defN tokenFile modeOption lthy =
let
val thy = Proof_Context.theory_of lthy
val ({src_path, lines, ...}:Token.file) = tokenFile
val path = if (Path.is_absolute src_path)
then src_path
else Path.append (Resources.master_directory thy) src_path
val mode = case modeOption of
SOME m => m
| NONE => "json"
val encoder = let
val _ = Convert_TensorFlow_Symtab.assert_target thy mode
in
case Convert_TensorFlow_Symtab.lookup_nn_encoder thy mode
of NONE => error "Encoder not found"
| SOME e => snd e
end
val strT = Nano_Json_Parser.stringT thy
val numT = Nano_Json_Parser.numT thy
val tf_json = Nano_Json_Parser.json_of_string (String.concat lines)
val json = TensorFlow_Json.transform_json (Path.dir path) tf_json
in
encoder defN strT numT json lthy
end
in
val _ = Outer_Syntax.command \<^command_keyword>‹import_TensorFlow›
"Import trained neural network from a TensorFlow.js JSON file."
((Parse.name -- \<^keyword>‹file› -- Resources.parse_file -- Scan.option (\<^keyword>‹as› |-- Parse.!!! (Parse.name))) >>
(fn (((name,_),get_file),mode_option) =>
Toplevel.theory (fn thy => let val file = get_file thy;
in Named_Target.theory_map (import_and_define_nn name file mode_option) thy end)
))
end
›
text‹
Lastly, we bind our encoder to a new top-level command: @{command import_TensorFlow} and prepare
its default configuration: ›
declare[[JSON_num_type = real, JSON_string_type = string, JSON_verbose = false]]
subsection‹Example Import›
text‹
In the following, we briefly demonstrate the use of the TensorFlow.js import.
›
import_TensorFlow compass file "examples/compass/model/model.json" as json
JSON_export compass file nor_model_transformed
text‹
This generated the definition @{const ‹compass›} with the following definition:
@{thm [display] "compass_def"}
›
end