Package: safejax
🔐 Serialize JAX, Flax, Haiku, or Objax model params withsafetensorssafejaxis a Python package to serialize JAX, Flax, Haiku, or Objax model params usingsafetensorsas the tensor storage format, instead of relying onpickle. For more details on whysafetensorsis safer thanpickleplease checkhuggingface/safetensors.Note thatsafejaxsupports the serialization ofjax,flax,dm-haiku, andobjaxmodel
parameters and has been tested with all those frameworks, but there may be some cases where it
does not work as expected, as this is still in an early development phase, so please if you have
any feedback or bug reports, open an issue atsafejax/issues.🛠️ Requirements & Installationsafejaxrequires Python 3.7 or abovepipinstallsafejax--upgrade💻 UsageflaxConvertparamstobytesin memoryfromsafejax.flaximportserialize,deserializeparams=model.init(...)encoded_bytes=serialize(params)decoded_params=deserialize(encoded_bytes)model.apply(decoded_params,...)Convertparamstobytesinparams.safetensorsfilefromsafejax.flaximportserialize,deserializeparams=model.init(...)encoded_bytes=serialize(params,filename="./params.safetensors")decoded_params=deserialize("./params.safetensors")model.apply(decoded_params,...)dm-haikuJust containsparamsfromsafejax.haikuimportserialize,deserializeparams=model.init(...)encoded_bytes=serialize(params)decoded_params=deserialize(encoded_bytes)model.apply(decoded_params,...)If it containsparamsandstatee.g. ExponentialMovingAverage in BatchNormfromsafejax.haikuimportserialize,deserializeparams,state=model.init(...)params_state={"params":params,"state":state}encoded_bytes=serialize(params_state)decoded_params_state=deserialize(encoded_bytes)# .keys() contains `params` and `state`model.apply(decoded_params_state["params"],decoded_params_state["state"],...)If it containsparamsandstate, but we want to serialize those individuallyfromsafejax.haikuimportserialize,deserializeparams,state=model.init(...)encoded_bytes=serialize(params)decoded_params=deserialize(encoded_bytes)encoded_bytes=serialize(state)decoded_state=deserialize(encoded_bytes)model.apply(decoded_params,decoded_state,...)objaxConvertparamstobytesin memory, and convert back toVarCollectionfromsafejax.objaximportserialize,deserializeparams=model.vars()encoded_bytes=serialize(params=params)decoded_params=deserialize(encoded_bytes)forkey,valueindecoded_params.items():ifkeyinmodel.vars():model.vars()[key].assign(value.value)model(...)Convertparamstobytesinparams.safetensorsfilefromsafejax.objaximportserialize,deserializeparams=model.vars()encoded_bytes=serialize(params=params,filename="./params.safetensors")decoded_params=deserialize("./params.safetensors")forkey,valueindecoded_params.items():ifkeyinmodel.vars():model.vars()[key].assign(value.value)model(...)Convertparamstobytesinparams.safetensorsand assign during deserializationfromsafejax.objaximportserialize,deserialize_with_assignmentparams=model.vars()encoded_bytes=serialize(params=params,filename="./params.safetensors")deserialize_with_assignment(filename="./params.safetensors",model_vars=params)model(...)More in-detail examples can be found atexamples/forflax,dm-haiku, andobjax.🤔 Whysafejax?safetensorsdefines an easy and fast (zero-copy) format to store tensors,
whilepicklehas some known weaknesses and security issues.safetensorsis also a storage format that is intended to be trivial to the framework
used to load the tensors. More in-depth information can be found athuggingface/safetensors.jaxusespytreesto store the model parameters in memory, so
it's a dictionary-like class containing nestedjnp.DeviceArraytensors.dm-haikuuses a custom dictionary formatted as<level_1>/~/<level_2>, where the
levels are the ones that define the tree structure and/~/is the separator between those
e.g.res_net50/~/intial_conv, and that key does not contain ajnp.DeviceArray, but a
dictionary with key value pairs e.g. for both weights aswand biases asb.objaxdefines a custom dictionary-like class namedVarCollectionthat contains
some variables inheriting fromBaseVarwhich is another customobjaxtype.flaxdefines a dictionary-like class namedFrozenDictthat is used to
store the tensors in memory, it can be dumped either intobytesinMessagePackformat or as astate_dict.Of all those,flaxis the only framework that defines its custom functions to
serialize and deserialize the model params underflax.serialization.Butflaxstill
usespickleas the format for storing the tensors, and there are no plans from HuggingFace
to extendsafetensorsto support anything more than tensors e.g.FrozenDicts, see their
response athuggingface/safetensors/discussions/138.So the motivation to createsafejaxis to easily provide a way to serializeFrozenDicts
usingsafetensorsas the tensor storage format instead ofpickle, as well as to provide
a common and easy way to serialize and deserialize any JAX model params (Flax, Haiku, or Objax)
usingsafetensorsformat.📄 Main differences withflax.serializationflax.serialization.to_bytesusespickleas the tensor storage format, whilesafejax.serializeusessafetensorsflax.serialization.from_bytesrequires thetargetto be instantiated, whilesafejax.deserializejust needs the encoded bytes🏋🏼 BenchmarkBenchmarks are no longer running withhyperfine,
as most of the elapsed time is not during the actual serialization but in the imports and
the model parameter initialization. So we've refactored those to run with pure
Python code usingtime.perf_counterto measure the elapsed time in seconds.$pythonbenchmarks/resnet50.py
safejax(100runs):2.0974s
flax(100runs):4.8734sThis means that forResNet50,safejaxis x2.3 times faster thanflax.serializationwhen
it comes to serialization, also to restate the fact thatsafejaxstores the tensors withsafetensorswhileflaxsaves those withpickle.But if we usehyperfineas mentioned above, it needs
to be installed first, and thehatch/pyenvenvironment needs to be activated
first (or just install the requirements). But, due to the overhead of the script, the
elapsed time during the serialization will be minimal compared to the rest, so the overall
result won't reflect well enough the efficiency diff between both approaches, as above.$hyperfine--warmup2"python benchmarks/hyperfine/resnet50.py serialization_safejax""python benchmarks/hyperfine/resnet50.py serialization_flax"Benchmark1:pythonbenchmarks/hyperfine/resnet50.pyserialization_safejaxTime(mean±σ):1.778s±0.038s[User:3.345s,System:0.511s]Range(min…max):1.741s…1.877s10runsBenchmark2:pythonbenchmarks/hyperfine/resnet50.pyserialization_flaxTime(mean±σ):1.790s±0.011s[User:3.371s,System:0.478s]Range(min…max):1.771s…1.810s10runsSummary'python benchmarks/hyperfine/resnet50.py serialization_safejax'ran1.01±0.02timesfasterthan'python benchmarks/hyperfine/resnet50.py serialization_flax'As we can see the difference is almost not noticeable, since the benchmark is using a
2-tensor dictionary, which should be faster using any method. The main difference is on
thesafetensorsusage for the tensor storage instead ofpickle.