              all_loss[sess] = loss
          logger.debug(all_loss)
          return all_loss
  
      def get_timestamps(self, session=None):
          """ Read the timestamps from the TensorBoard logs.
  
          As loss timestamps are slightly different for each loss, we collect the timestamp from the
          `batch_total` key.
  
 
</s> remove         logger.debug("Getting loss: (session: %s)", session)
        all_loss = dict()
        for sess, logfile in self.log_filenames.items():
            if session is not None and sess != session:
                logger.debug("Skipping session: %s", sess)
                continue
            loss = dict()
            events = [event_pb2.Event.FromString(record.numpy())
                      for record in tf.data.TFRecordDataset(logfile)]
            for event in events:
                if not event.summary.value or not event.summary.value[0].tag.startswith("batch_"):
                    continue
                summary = event.summary.value[0]
                tag = summary.tag.replace("batch_", "")
                loss.setdefault(tag, []).append(summary.simple_value)
            all_loss[sess] = loss
        logger.debug(all_loss)
        return all_loss
 </s> add         logger.debug("Getting loss: (session_id: %s)", session_id)
        retval = dict()
        for sess, info in self._from_cache(session_id, is_training).items():
            arr = np.frombuffer(zlib.decompress(info["loss"]),
                                dtype="float32").reshape(info["loss_shape"])
            for idx, title in enumerate(info["labels"]):
                retval.setdefault(sess, dict())[title] = arr[:, idx]
        logger.debug({key: {k: v.shape for k, v in val.items()}
                      for key, val in retval.items()})
        return retval </s> remove     """ Parse and return data from TensorBoard logs """
 </s> add     """ Parse data from TensorBoard logs.

    Process the input logs folder and stores the individual filenames per session.

    Caches timestamp and loss data on request and returns this data from the cache.

    Parameters
    ----------
    logs_folder: str
        The folder that contains the Tensorboard log files
    """ </s> remove         session: int, optional
 </s> add         session_id: int, optional </s> remove     def get_loss(self, session=None):
 </s> add     def _cache_data(self, session_id, is_training=False):
        """ Cache TensorBoard logs for the given session ID on first access.

        Populates :attr:`_cache` with timestamps and loss data.

        If this is a training session and the data is being queried for the training session ID
        then get the latest available data and append to the cache

        Parameters
        -------
        session_id: int
            The session ID to cache the data for
        is_training: bool, optional
            ``True`` if a current training session is running otherwise ``False``.
            Default: ``False``
        """
        labels = []
        step = []
        loss = []
        timestamps = []
        last_step = -1
        live_data = is_training and session_id == max(self._log_filenames)

        if live_data:
            iterator = self._get_latest_live()
        else:
            iterator = tf.compat.v1.io.tf_record_iterator(self._log_filenames[session_id])

        try:
            for record in iterator:
                event = event_pb2.Event.FromString(record)
                if not event.summary.value or not event.summary.value[0].tag.startswith("batch_"):
                    continue

                if last_step == -1:
                    last_step = event.step if live_data else 0

                if event.step != last_step:
                    loss.append(step)
                    step = []
                    last_step = event.step

                summary = event.summary.value[0]
                tag = summary.tag

                if tag == "batch_total":
                    timestamps.append(event.wall_time)
                    continue

                lbl = tag.replace("batch_", "")
                if lbl not in labels:
                    labels.append(lbl)

                step.append(summary.simple_value)

        except tf_errors.DataLossError as err:
            logger.warning("The logs for Session %s are corrupted and cannot be displayed. "
                           "The totals do not include this session. Original error message: "
                           "'%s'", session_id, str(err))

        if step:
            loss.append(step)

        loss = np.array(loss, dtype="float32")
        timestamps = np.array(timestamps, dtype="float64")

        logger.debug("Caching session id: %s, labels: %s, loss: %s, timestamps: %s",
                     session_id, labels, loss.shape, timestamps.shape)

        if live_data and session_id in self._cache:
            self._add_latest_data_to_cache(session_id, loss, timestamps)
        else:
            self._cache[session_id] = dict(labels=labels,
                                           loss=zlib.compress(loss),
                                           loss_shape=loss.shape,
                                           timestamps=zlib.compress(timestamps),
                                           timestamps_shape=timestamps.shape)

    def _get_latest_live(self):
        """ Obtain the latest event logs for live training data and add to the cache """
        if self._training_iterator is None:
            training_session_id = self.session_ids[-1]
            filename = self._log_filenames[training_session_id]
            self._training_iterator = tf.compat.v1.io.tf_record_iterator(filename)
            logger.debug("Set live training iterator %s for session_id: %s",
                         self._training_iterator, training_session_id)

        i = 0
        while True:
            try:
                yield next(self._training_iterator)
                i += 1
            except StopIteration:
                logger.debug("End of data reached")
                break
            except tf.errors.DataLossError as err:
                # Truncated records are ignored. The iterator holds the offset, so the record will
                # be completed at the next call.
                logger.debug("Truncated record. Original Error: %s", err)
                break
        logger.debug("Collected %s records from live log file", i)

    def _add_latest_data_to_cache(self, session_id, loss, timestamps):
        """ Append the latest received live training data to the cached data.

        Parameters
        ----------
        session_id: int
            The training session ID to update the cache for
        loss: :class:`numpy.ndarray`
            The latest loss values returned from the iterator
        timestamps: :class:`numpy.ndarray`
            The latest time stamps returned from the iterator
        """
        if not np.any(loss) and not np.any(timestamps):
            logger.debug("No new live data to cache.")
            return

        logger.debug("Adding live data to cache: (loss: %s, timestamps: %s)",
                     loss.shape, timestamps.shape)

        cache = self._cache[session_id]

        past_loss = np.frombuffer(zlib.decompress(cache["loss"]),
                                  dtype="float32").reshape(cache["loss_shape"])
        new_loss = np.concatenate((past_loss, loss))
        cache["loss_shape"] = new_loss.shape
        cache["loss"] = zlib.compress(new_loss)
        del past_loss

        past_timestamps = np.frombuffer(zlib.decompress(cache["timestamps"]),
                                        dtype="float64").reshape(cache["timestamps_shape"])
        new_timestamps = np.concatenate((past_timestamps, timestamps))
        cache["timestamps_shape"] = new_timestamps.shape
        cache["timestamps"] = zlib.compress(new_timestamps)
        del past_timestamps

    def _from_cache(self, session_id=None, is_training=False):
        """ Get the session data from the cache.

        If the request data does not exist in the cache, then populate it.

        Parameters
        ----------
        session_id: int, optional
            The Session ID to return the data for. Set to ``None`` to return all session
            data. Default ``None`
        is_training: bool, optional
            ``True`` if a current training session is running otherwise ``False``.
            Default: ``False``

        Returns
        -------
        dict
            The session id(s) as key, with the event data as value
        """
        if session_id is not None and session_id not in self._cache:
            self._cache_data(session_id)
        elif is_training and session_id == self.session_ids[-1]:
            self._cache_data(session_id, is_training=is_training)
        elif session_id is None and not all(idx in self._cache for idx in self._log_filenames):
            for sess in self._log_filenames:
                if sess not in self._cache:
                    self._cache_data(sess)

        if session_id is None:
            return self._cache
        return {session_id: self._cache[session_id]}

    def get_loss(self, session_id=None, is_training=False): </s> remove         session: int, optional
 </s> add         session_id: int, optional
    def get_timestamps(self, session_id=None, is_training=False):
              all_loss[sess] = loss
          logger.debug(all_loss)
          return all_loss
  
    def get_timestamps(self, session_id=None, is_training=False):
          """ Read the timestamps from the TensorBoard logs.
  
          As loss timestamps are slightly different for each loss, we collect the timestamp from the
          `batch_total` key.
  
keep keep keep keep replace keep keep keep keep keep
 GUI - Stats optimization (#1067)

* Faster stat loading + caching

* Compress data in cache

* Optimize some calculations

* Vectorize smoothing

* stats.Calculations optimized

* Load latest training data from live iterator

* Add options to training graph 
lib/gui/stats.py