tfmanager

Overview

class tfmanager.TFManager(graph_info, device, q, positions_buffer, nlist_buffer, forces_buffer, box_buffer, virial_buffer, log_filename, dtype, debug, write_tensorboard, use_feed, bootstrap, primary, bootstrap_map, save_period, use_xla, mol_indices)

TFManager manages TensorFlow training and coordinates communication with HOOMD. Sets up the TensorFlow graph, sets placeholder variable values, updates and saves the TensorFlow graph, passes output forces back to HOOMD, and writes TensorBoard files.

Parameters:
  • graph_info – The structure of the TensorFlow graph, passed as a dict. See tfcompute.py
  • device – Which device to run on. See tfcompute.py
  • q – Threading queue that is used during execution of TensorFlow.
  • positions_buffer – Buffer where particle positions are stored
  • nlist_buffer – Address of the neighbor list tensor
  • box_buffer – Address of the box tensor
  • forces_buffer – Address of the forces tensor
  • virial_buffer – Address of the virial tensor
  • log_filename – Name of the file to output tensorflow logs
  • dtype – Data type for tensor values, e.g. int32, float32, etc
  • debug – True to run TensorFlow in debug mode
  • write_tensorboard – Whether to output a tensorboard file
  • use_feed – Whether or not to use a feed dictionary of tensor values
  • bootstrap – Location of previously-trained model files to load, otherwise None
  • primary – Whether this is the ‘primary’ instance of TFManager. Only one instance writes logs and saves model files.
  • bootstrap_map – A dictionary to be used when bootstrapping, pairing old models’ tensor variable names with new ones. Key is new name, value is older model’s.
  • save_period – How often to save the TensorFlow data. Period here is measured by how many times the TensorFLow model is updated. See tfcompute.py.
  • use_xla – If True, enables the accelerated linear algebra library in TensorFlow, which can be useful for large and complicated tensor operations.
  • mol_indices – The molecucule indices, if doing mol batch (None otherwise)
start_loop()

start_loop method of tfmanager. Prepares GPU for execution, gathers trainable variables, sets up model saving, loads pre-trained variables, sets up tensorboard if requested and parses feed_dicts.