itwinai.tensorflow

distributed

TensorFlow distributed strategies.

itwinai.tensorflow.distributed.get_strategy() → Tuple[tensorflow.distribute.Strategy, int][source]

Strategy for distributed TensorFlow training. It will automatically detect if you are running in a multi-node environment, returning the correct TensorFlow distributed strategy for data parallel distributed training.

Returns:

strategy and number of parallel workers.

Return type:

Tuple[tf.distribute.Strategy, int]

trainer

Base TensorFlow trainer module.

class itwinai.tensorflow.trainer.TensorflowTrainer(epochs: int, micro_batch_size: int, shuffle_buffer: int | None = None, callbacks: List[Dict | keras.callbacks.Callback] | None = None, model_config: Dict | None = None, model_compile_config: Dict | None = None, rnd_seed: int | None = None, verbose: str | int = 'auto')[source]

Bases: Trainer

Trains a Keras model.

Parameters:

  • epochs (int) – number of training epochs.

  • micro_batch_size (int) – per-worker batch size. Equals macro batch size when not running distributed.

  • shuffle_buffer (Optional[int], optional) – if given, shuffles dataset using a buffer of given size. See tf.data.Dataset.shuffle. Defaults to None.

  • callbacks (Optional[List], optional) – list fo Keras callbacks. Can be a list of dictionary configurations. Defaults to None.

  • model_config (Optional[Dict], optional) – model configuration. If given, a model is instantiated from this configuration. Defaults to None.

  • model_compile_config (Optional[Dict], optional) – configuration for keras.Model.compile. Defaults to None.

  • rnd_seed (Optional[int], optional) – random seed. Defaults to None.

  • verbose (Union[str, int], optional) – verbosity level for keras.Model.fit. Defaults to β€˜auto’.

epochs: int

Total number of training epochs.

micro_batch_size: int

Per-worker batch size (when distributed).

shuffle_buffer: int | None = None

Buffer used to shuffle dataset. Defaults to None.

rnd_seed: int | None = None

Random seed for reproducibility. Defaults to None.

callbacks: List | None = None

List of Keras callbacks. Defaults to None.

strategy: tensorflow.distribute.Strategy

TensorFlow distributed strategy.

num_workers: int

Total number of workers in distributed strategy.

macro_batch_size: int

Total batch size. When distributed, it is the sum of micro_batch_size across all workers.

static instantiate_compile_conf(model_compile_config: Dict) → Dict[str, Any][source]

Instantiate fields of Keras model.compile() from their dictionary serialization.

Parameters:

model_compile_config (Dict) – fields of Keras model.compile() serialized as dictionary.

Returns:

dictionary mapping compile argument names to the instantiated objects.

Return type:

Dict[str, Any]

static instantiate_callbacks(callbacks: List[Dict | keras.callbacks.Callback]) → List[keras.callbacks.Callback][source]

Instantiate Keras callbacks from dictionaries.

Parameters:

  • callbacks (List[Union[Dict, Callback]]) – list of Keras callbacks

  • dictionary. (in serialized as)

Returns:

list of instantiated callbacks.

Return type:

List[Callback]

execute(train_dataset: tensorflow.data.Dataset, validation_dataset: tensorflow.data.Dataset, test_dataset: tensorflow.data.Dataset | None = None) → Tuple[tensorflow.data.Dataset, tensorflow.data.Dataset, tensorflow.data.Dataset, keras.Model][source]

Run training. Users should override this method.

Parameters:

  • train_dataset (Dataset) – train dataset of type tensorflow.data.Dataset.

  • validation_dataset (Dataset) – validation dataset of type tensorflow.data.Dataset.

  • test_dataset (Optional[Dataset], optional) – test dataset of type tensorflow.data.Dataset. Defaults to None.

Returns:

tuple of train_dataset, validation_dataset, test_dataset, and trained Keras model.

Return type:

Tuple[Dataset, Dataset, Dataset, keras.Model]

utils

itwinai.tensorflow.utils.model_to_json(model: keras.Model, filepath: str)[source]

Serialize Keras model to JSON file.

Parameters:

  • model (keras.Model) – Keras model.

  • filepath (str) – JSON file path.

itwinai.tensorflow.utils.model_from_json(filepath: str) → keras.Model[source]

Deserialize Keras model from JSON file.

Parameters:

filepath (str) – JSON file path.

Returns:

loaded Keras model.

Return type:

keras.Model