Utilities

Cache

Loading pretrained weights for models via

create_model(..., pretrained=True)

usually requires downloading the weights from the internet. To avoid that, either to make model creation faster or when executing code in environments without internet access, we provide a two-part cache system.

There is a global cache directory, accessible by get_dir(). The cache directory is determined in the following order of precedence

• Directory set by set_dir()

• Path set by env variable $TFIMM_HOME

• Path $XDG_CACHE_HOME/tfimm if env variable $XDG_CACHE_HOME is set

• ~/.cache

Within the cache directory each model should be saved in a subdirectory with name model_name. If the cache directory is ~/.cache, then resnet18 should be saved under ~/.cache/resnet18.

In some cases it might be impractical to store all models in the same directory. In that case we can specify the cache location for individual models using set_model_cache(). For example, if we have resnet18 weights stored at ~/resnet18_saved, we call

set_model_cache("resnet18", "~/resnet18_saved")
# This will read weights from the above location
model = create_model("resnet18", pretrained=True)

We can remove the cache location for a particular model by calling clear_model_cache(). This will only stop tfimm from looking in that directory, it will not delete anything from disk.

We can also list all models for which we have specified locations via list_cached_models(). This will return a list of all models whose cache location has been set via set_model_cache(). It will not look in the general cache directory.

We can obtain the cache location for a particular model via cached_model_path(). The cache location is determined in the following order of precedence:

• Model-specific cache set via set_model_cache()

• General cache directory given by get_dir()

cached_model_path(model_name)

Checks if the weights for model model_name are cached. If so, we return the path to the weights, otherwise return None. We check the model-speficif cache first before looking in the general cache directory.

Parameters:

model_name (str) – Model to be queried.

Return type:

str | None

clear_model_cache(model_name)

Removes the cache entry for model_name.

Parameters:

model_name (str) – Model name. Same name that is passed to create_model().

get_dir()

Get the tfimm cache directory used for storing downloaded models & weights.

If tfimm.set_dir has not been called, default path is $TFIMM_HOME, where environment variable $TFIMM_HOME defaults to $XDG_CACHE_HOME/tfimm. $XDG_CACHE_HOME follows the X Design Group specifications of the Linux filesystem layout, with a default value ~/.cache if the environment variable is not set.

Return type:

str

list_cached_models()

Lists all models in model-specific cache.

Returns:

List of model names.

Return type:

List[str]

set_dir(d)

Optionally set the tfimm models directory used to save downloaded models & weights.

Parameters:

d (str) – Path to a local folder to save downloaded models & weights.

set_model_cache(model_name, model_path)

Sets the cache location for model_name to model_path.

Parameters:

• model_name (str) – Model name. Same name that is passed to create_model().

• model_path (str) – Path, where the model is stored.