Function main

examples/mnist_inference.py:32–60

()

def main():
    mnist = tf.keras.datasets.mnist
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0

    tracker.start()

    # First mode of task emission tracking, using explicit functions start_task & stop_task
    tracker.start_task("build model")
    model = KerasClassifier(build_fn=build_model, epochs=1)
    tracker.stop_task()
    param_grid = dict(batch_size=list(range(32, 256 + 32, 32)))

    # Track task emissions using the context manager
    with TaskEmissionsTracker(task_name="Grid search", tracker=tracker):
        grid = GridSearchCV(estimator=model, param_grid=param_grid)
    grid.fit(x_train, y_train)

    for _ in range(10):
        # Third tracking mode for tasks, use a decorated function with track_task_emissions decorator
        predict(grid, x_test)

    emissions = tracker.stop()

    print(f"Emissions : {emissions} kg CO₂")
    for task_name, task in tracker._tasks.items():
        print(
            f"Emissions : {task.emissions_data.emissions} kg CO₂ for task {task_name}"
        )


if __name__ == "__main__":