Skip to content

XVIII. Bioimage.io III (load models)

carlosuc3m edited this page Sep 13, 2023 · 2 revisions

JDLL was first developed to load models from the Bioimage.io in the Icy Bioimaging software and although it evolved into a library that can be used to build more general software, it is still tightly integrated and linked to the Bioimage.io.

The thigh integration with the Bioimage.io model repo allows retrieving the models available there and their information and specifications, an easy download of Bioimage.io models, an easy load of Bioimage.io models and methods to retrieve the specifications of the Bioimage.io models.

Load Bioimage.io models

The integration of the Bioiamge.io into JDLL easies a lot the process of loading models to be able to make inference with them.

The standard loading process implies several steps to define which of the Deep Learing frameworks (engines) needs to be used and some parameters such as the modelSource as defined in the weigths field of the [rdf.yaml file](https://github.com/bioimage-io/spec-bioimage-io/blob/gh-pages/model_spec_latest.md. However, for Bioimage.io models, all the technical details needed to load the model are contained in the rdf.yaml file that is located inside of the model folder.

Bioimage.io models can be loaded easily with just one command:

// Folder containing the model. This is the parent directory
// of the rdf.yaml file that also contains the weights and other files
String modelFolder = "/path/to/model";
// Directory where all the engines folders are located
String enginesFolder = "/path/to/engines";
// Create the model
Model model = Model.createBioimageioModel(modelFolder , enginesFolder);
// Load the model in the memory
model.load();

Note that if the needed Deep Learning framework (DL framework for which the model has been trained) has not been previously installed, the model will not be created and exception will be thrown by the line:

Model model = Model.createBioimageioModel(modelFolder , enginesFolder);

The information about the engines the model was trained for is contained in the rdf.yaml file and can be retreived creating an instance of ModelDescriptor:

String localRdfYaml = "/path/to/model/rdf.yaml";
boolean verbose = false;
// Create the model descriptor
ModelDescriptor descriptor = ModelDescriptor.readFromLocalFile(localRdfYaml, verbose);
// Retreive the weight DL formats in which the model was trained
ModelWeight weights = descriptor.getWeights();
for (WeightFormat ww : weights.getSupportedWeights())
  System.out.println(ww.getWeightsFormat() + " " + ww.getTrainingVersion());

For more info, about the ModelDescriptor explore this section of the wiki.

If any of the required engines is not installed, install it to be able to load the model. For more information about engine management and installation click here.

However, regard that

Model model = Model.createBioimageioModel(modelFolder , enginesFolder);

will be able to load the model as long as a compatible engine version with the model is installed. Compatible versions are defined as engines from the same Deep Learning framework that share the same major version (first number in the version string). For example, Pytorch 1.13.0 and 1.9.1 are compatible, but Tensorflow 1.15.0 and Tensorflow 2.7.0 are not.

For more information about version compatibility, click here.

Methods to create models

REgard that the following methods only work for Bioimage.io models. This is, models that have been downloaded from the Bioiamge.io and that contain all the files required for a Bioimage.io model (at least the weigths file and the rdf.yaml file).

Model.createBioimageioModel(String bmzModelFolder, String enginesFolder)

Method that creates a Bioimage.io model defined by the path referred by the bmzModelFolder arg. The method will look in the engines directory defined by enginesFolder to identify if any of the installed engines will be able to load the model weights.

Regard that this method allows compatible versions, which means that if the weights were trained by (for example) Tensorflow 2.4.0, Tensorflow 2.7.0 will be allowed to try to to load it. Also note that this might induce some errors if any of the supported methods of Tensorflow 2.4.0 is not supported in Tensorflow 2.7.0 and it is used in the model of interest. For more information about compatible versions click here.

  • bmzModelFolder: path to the Bioimage.io model folder
  • enginesFolder: path to the directory where the JDLL Deep Learning engines have been installed

Model.createBioimageioModelWithExactWeigths(String bmzModelFolder, String enginesFolder)

Method that creates a Bioimage.io model defined by the path referred by the bmzModelFolder arg. The method will look in the engines directory defined by enginesFolder to identify if any of the installed engines will be able to load the model weights.

Regard that this method only creates the model if the DL engine version is the same as the version of the weigths that were used to train the model, which is specified in the model rdf.yaml file.

Even though this method might increase the complexity and increase the amunt of engines downloaded by the user, it is safer as it makes sure that there will be no conflicts when trying to load and run the model. For more info about checking if the engines for a Bioimage.io model are installed click here.

  • bmzModelFolder: path to the Bioimage.io model folder
  • enginesFolder: path to the directory where the JDLL Deep Learning engines have been installed

Providing a classloader

In order to load a model and avoid conflicts with other engines, JDLL creates a separate classloader that loads the JARs needed to run the corresponding DL framework. The new classloader is and URLClassLoader using the context classloader (Thread.currentThread().getContextClassLoader()) as its parent.

However, there are some examples in which the application or software using JDLL might have designed their custom approach to handle classloading. For those cases the methods used to create Deep Learning models can digest one argument more at the end: ClassLoader classLoader.

The new call to the method would be the same as in the begining of the section but with one more arg:

Model model = Model.createBioimageioModel(modelFolder , enginesFolder, classloader);

If classloader=null, the method will fallback to the context classloader (as if the argument had not been provided): Thread.currentThread().getContextClassLoader().

Check if the engine for a model is installed

As explained in the previous section it necessary to know which are the engines installed and which are the engines required to load a model.

This section will explain how to check if any of the engines required to load a model is installed or not.

First, we need to retrieve the weight formats available in the Bioimage.io model. The weight formats refer to the DL framework and version in which the models were trained. Thus to load a model in certain weight format and version, we need to have installed that engine of that version.

There are several weight formats available in the Bioimage.io: keras_hdf5, onnx, pytorch_state_dict, tensorflow_js, tensorflow_saved_model_bundle, and torchscript.

Both pytorch_state_dict and torchscript are used to load Pytorch models, however JDLL is only compatible with torchscript as it is the only one that is compatible with Java. Same with tensorflow_js and tensorflow_saved_model_bundle, only tensorflow_saved_model_bundle is compatible with Java. To see which are the supported engines by JDLL click here.

To retrieve the specifications and information about a model from its rdf.yaml file:

String rdfPath = "/path/to/model/rdf.yaml":
boolean verbose = false;
// Create the model descriptor
ModelDescriptor descriptor = ModelDescriptor.readFromLocalFile(rdfPath , verbose);
// Retreive the weight DL formats in which the model was trained
ModelWeight weights = descriptor.getWeights();

Once we have the information about the model and about the weights we can use the method AvailableEngines.isEngineSupported(String framework, String version, Boolean cpu, Boolean gpu). This method filters the engines supported by JDLL and returns true if at least one of them matches our search.

Following our example, if we want to check if the first weight format (the weights of a Bioimage.io model can be in one or more formats) supported by the model is allowed by JDLL we can do:

WeightFormat firstFormat = weights.getSupportedWeights().get(0);
// We can set the parameters 'cpu' and 'gpu' to null because we do not care
// about cpu or gpu support right now
boolean supported = 
                AvailableEngines.isEngineSupported(firstFormat.getWeightsFormat(), 
                                                    firstFormat.getTrainingVersion(), null, null);

If the engine is installed, we can now check whether it has been already installed. The method used is InstalledEngines.checkEngineWithArgsInstalled(String engine, String version, Boolean cpu, Boolean gpu, Boolean rosetta, String enginesDir), where rosetta is only useful for MacOS M (see here) and enginesDir refers to the directory where all the engines are installed:

String framework = firstFormat.getWeightsFormat();
String version = firstFormat.getTrainingVersion();
String enginesDir = "/path/to/engines/dir";
List<DeepLearningVersion> versions = InstalledEngines.checkEngineWithArgsInstalled(framework , version, null, null, enginesDir);
boolean installed = versions.size() > 0;

If installed == true we can then proceede to create and load the model.

This code will also work for the same purpose:

String engine = "tensorflow";
String version = "1.13.1";
String enginesDir = "/path/to/engines";
boolean installed = InstalledEngines.checkEngineVersionInstalled(engine, version, enginesDir);

If installed == true we can then proceede to create and load the model.

Notice that in the previous case we are checking for the same exact version as that one in the weights, this would be the procedure to check if we can create and load a model using Model.createBioimageioModelWithExactWeigths(String bmzModelFolder, String enginesFolder).

However, as mentioned above, we do not need the exact same version to load an engine, just a compatible one, if we use the method Model.createBioimageioModel(String bmzModelFolder, String enginesFolder).

To check if there exist any compatible version to the one required by the model weights we can use the method InstalledEngines.getMostCompatibleVersionForEngine(String engine, String version, String enginesDir) which will return null if there is no compatible engine or a String containing the version of the most compatible version to the engine version of interest:

String framework = firstFormat.getWeightsFormat();
String version = firstFormat.getTrainingVersion();
String enginesDir = "/path/to/engines/dir";
String compat = InstalledEngines.getMostCompatibleVersionForEngine(framework, version, String enginesDir);
boolean installed = compat != null;

If installed == true we can then proceede to create and load the model with Model.createBioimageioModel(String bmzModelFolder, String enginesFolder).