|
1 | 1 | # coding: utf-8
|
2 |
| -"""Trains and evaluate a simple MLP |
3 |
| -on the Reuters newswire topic classification task. |
4 |
| -""" |
5 |
| -from __future__ import print_function |
| 2 | +import os |
6 | 3 |
|
7 |
| -# The following imports are the only additions to code required |
8 |
| -# to automatically log metrics and parameters to Comet. |
9 |
| -import comet_ml # noqa |
| 4 | +import comet_ml |
10 | 5 |
|
11 |
| -import mlflow.keras |
12 |
| -import numpy as np |
| 6 | +# You can use 'tensorflow', 'torch' or 'jax' as backend. Make sure to set the |
| 7 | +# environment variable before importing. |
| 8 | +os.environ["KERAS_BACKEND"] = "tensorflow" |
13 | 9 |
|
14 |
| -import keras |
15 | 10 |
|
16 |
| -# The following import and function call are the only additions to code required |
17 |
| -# to automatically log metrics and parameters to MLflow. |
18 |
| -import mlflow |
19 |
| -from keras.datasets import reuters |
20 |
| -from keras.layers import Activation, Dense, Dropout |
21 |
| -from keras.models import Sequential |
22 |
| -from keras.preprocessing.text import Tokenizer |
| 11 | +import mlflow.keras # noqa: E402 |
| 12 | +import numpy as np # noqa: E402 |
23 | 13 |
|
24 |
| -# The sqlite store is needed for the model registry |
25 |
| -mlflow.set_tracking_uri("sqlite:///db.sqlite") |
| 14 | +import keras # noqa: E402 |
26 | 15 |
|
27 |
| -# We need to create a run before calling keras or MLFlow will end the run by itself |
28 |
| -mlflow.set_experiment("comet-example-mlflow-hello-world") |
29 |
| -mlflow.start_run() |
| 16 | +# Login to Comet if necessary |
| 17 | +comet_ml.login(project_name="comet-example-mlflow-hello-world") |
30 | 18 |
|
31 |
| -mlflow.keras.autolog() |
| 19 | +# Load dataset |
| 20 | +(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data() |
| 21 | +x_train = np.expand_dims(x_train, axis=3) |
| 22 | +x_test = np.expand_dims(x_test, axis=3) |
| 23 | +x_train[0].shape |
32 | 24 |
|
33 |
| -max_words = 1000 |
34 |
| -batch_size = 32 |
35 |
| -epochs = 5 |
| 25 | +# Build model |
| 26 | +NUM_CLASSES = 10 |
| 27 | +INPUT_SHAPE = (28, 28, 1) |
| 28 | + |
| 29 | + |
| 30 | +def initialize_model(): |
| 31 | + return keras.Sequential( |
| 32 | + [ |
| 33 | + keras.Input(shape=INPUT_SHAPE), |
| 34 | + keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"), |
| 35 | + keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"), |
| 36 | + keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"), |
| 37 | + keras.layers.GlobalAveragePooling2D(), |
| 38 | + keras.layers.Dense(NUM_CLASSES, activation="softmax"), |
| 39 | + ] |
| 40 | + ) |
36 | 41 |
|
37 |
| -print("Loading data...") |
38 |
| -(x_train, y_train), (x_test, y_test) = reuters.load_data( |
39 |
| - num_words=max_words, test_split=0.2 |
40 |
| -) |
41 | 42 |
|
42 |
| -print(len(x_train), "train sequences") |
43 |
| -print(len(x_test), "test sequences") |
| 43 | +model = initialize_model() |
| 44 | +model.summary() |
44 | 45 |
|
45 |
| -num_classes = np.max(y_train) + 1 |
46 |
| -print(num_classes, "classes") |
| 46 | +# Train model |
47 | 47 |
|
48 |
| -print("Vectorizing sequence data...") |
49 |
| -tokenizer = Tokenizer(num_words=max_words) |
50 |
| -x_train = tokenizer.sequences_to_matrix(x_train, mode="binary") |
51 |
| -x_test = tokenizer.sequences_to_matrix(x_test, mode="binary") |
52 |
| -print("x_train shape:", x_train.shape) |
53 |
| -print("x_test shape:", x_test.shape) |
| 48 | +BATCH_SIZE = 64 # adjust this based on the memory of your machine |
| 49 | +EPOCHS = 3 |
54 | 50 |
|
55 |
| -print( |
56 |
| - "Convert class vector to binary class matrix " |
57 |
| - "(for use with categorical_crossentropy)" |
| 51 | +model = initialize_model() |
| 52 | + |
| 53 | +model.compile( |
| 54 | + loss=keras.losses.SparseCategoricalCrossentropy(), |
| 55 | + optimizer=keras.optimizers.Adam(), |
| 56 | + metrics=["accuracy"], |
58 | 57 | )
|
59 |
| -y_train = keras.utils.to_categorical(y_train, num_classes) |
60 |
| -y_test = keras.utils.to_categorical(y_test, num_classes) |
61 |
| -print("y_train shape:", y_train.shape) |
62 |
| -print("y_test shape:", y_test.shape) |
63 |
| - |
64 |
| -print("Building model...") |
65 |
| -model = Sequential() |
66 |
| -model.add(Dense(512, input_shape=(max_words,))) |
67 |
| -model.add(Activation("relu")) |
68 |
| -model.add(Dropout(0.5)) |
69 |
| -model.add(Dense(num_classes)) |
70 |
| -model.add(Activation("softmax")) |
71 |
| - |
72 |
| -model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]) |
73 |
| - |
74 |
| -history = model.fit( |
| 58 | + |
| 59 | +run = mlflow.start_run() |
| 60 | +model.fit( |
75 | 61 | x_train,
|
76 | 62 | y_train,
|
77 |
| - batch_size=batch_size, |
78 |
| - epochs=epochs, |
79 |
| - verbose=1, |
| 63 | + batch_size=BATCH_SIZE, |
| 64 | + epochs=EPOCHS, |
80 | 65 | validation_split=0.1,
|
| 66 | + callbacks=[mlflow.keras.MlflowCallback(run)], |
81 | 67 | )
|
82 |
| -score = model.evaluate(x_test, y_test, batch_size=batch_size, verbose=1) |
83 |
| -print("Test score:", score[0]) |
84 |
| -print("Test accuracy:", score[1]) |
85 | 68 |
|
86 | 69 | mlflow.keras.log_model(model, "model", registered_model_name="Test Model")
|
| 70 | + |
| 71 | +mlflow.end_run() |
0 commit comments