def EpochTelemetry(self) -> typing.List[telemetry_pb2.ModelEpochTelemetry]:
"""Return the epoch telemetry files."""
return [
pbutil.FromFile(self.logdir / p, telemetry_pb2.ModelEpochTelemetry())
for p in sorted(self.logdir.iterdir())
if re.match(r"epoch_\d\d+_telemetry\.pbtxt", str(p.name))
]
def EpochEndCallback(self, epoch: int, loss: float):
now = labdate.MillisecondsTimestamp()
epoch_time_ms = now - self.last_epoch_begin_timestamp
telemetry = telemetry_pb2.ModelEpochTelemetry(
timestamp_utc_epoch_ms=now,
epoch_num=epoch,
epoch_wall_time_ms=epoch_time_ms,
loss=loss,
)
pbutil.ToFile(telemetry, self.logdir / f"epoch_{epoch:03d}_telemetry.pbtxt")