wehub-resource-sync
25 lines
790 B
Python
25 lines
790 B
Python
import time
|
|
|
|
import keras
|
|
|
|
|
|
class BenchmarkMetricsCallback(keras.callbacks.Callback):
|
|
def __init__(self, start_batch=1, stop_batch=None):
|
|
self.start_batch = start_batch
|
|
self.stop_batch = stop_batch
|
|
|
|
# Store the throughput of each epoch.
|
|
self.state = {"throughput": []}
|
|
|
|
def on_train_batch_begin(self, batch, logs=None):
|
|
if batch == self.start_batch:
|
|
self.state["epoch_begin_time"] = time.time()
|
|
|
|
def on_train_batch_end(self, batch, logs=None):
|
|
if batch == self.stop_batch:
|
|
epoch_end_time = time.time()
|
|
throughput = (self.stop_batch - self.start_batch + 1) / (
|
|
epoch_end_time - self.state["epoch_begin_time"]
|
|
)
|
|
self.state["throughput"].append(throughput)
|