def do_train(args):
model = Net()
if args.cuda:
model.cuda()
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
train_loader = utils.get_data(True, args.batch_size)
test_loader = utils.get_data(False, args.batch_size)
with mlflow.start_run() as run:
print("run_id:",run.info.run_id)
mlflow.set_tag("mlflow_version",mlflow.__version__)
mlflow.set_tag("torch_version",torch.__version__)
mlflow.set_tag("autolog",args.autolog)
# Perform the training
for epoch in range(1, args.epochs + 1):
train(model, optimizer, epoch, args, train_loader)
test(model, epoch, args, train_loader, test_loader)
if not args.autolog:
for key, value in vars(args).items():
mlflow.log_param(key, value)
mlflow.pytorch.log_model(model, "pytorch-model")
# Log model as ONNX
if args.log_as_onnx:
import onnx_utils
import onnx
print("ONNX Version:", onnx.__version__)
dataiter = iter(test_loader)
images, labels = dataiter.next()
onnx_utils.log_model(model, "onnx-model", images)
def run(epochs, log_as_onnx):
model = Model()
print("model.type:",type(model))
criterion = torch.nn.MSELoss(size_average=False)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
print("Train:")
for epoch in range(epochs):
y_pred = model(x_data)
loss = criterion(y_pred, y_data)
print(f" Epoch: {epoch} Loss: {loss.data.item()}")
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Predictions:")
for v in test_data:
tv = torch.Tensor([[v]])
y_pred = model(tv)
print(f" {v}: {model(tv).data[0][0]}")
with mlflow.start_run() as run:
print("run_id:",run.info.run_id)
mlflow.log_param("epochs", epochs)
mlflow.pytorch.log_model(model, "pytorch-model")
if args.log_as_onnx:
import onnx_utils
import onnx
print("ONNX Version:", onnx.__version__)
onnx_utils.log_model(model, "onnx-model", x_data)
def train(run, model_name, data_path, epochs, batch_size, mlflow_custom_log, log_as_onnx):
x_train, y_train, x_test, y_test = utils.get_train_data(data_path)
model = build_model()
model.compile(
optimizer="rmsprop",
loss="categorical_crossentropy",
metrics=["accuracy"])
model.summary()
model.fit(x_train, y_train, epochs=epochs, batch_size=batch_size, verbose=1, callbacks=[LogMetricsCallback()])
print("model.type:", type(model))
test_loss, test_acc = model.evaluate(x_test, y_test)
print("test_acc:", test_acc)
print("test_loss:", test_loss)
if mlflow_custom_log:
mlflow.log_param("epochs", epochs)
mlflow.log_param("batch_size", batch_size)
mlflow.log_metric("test_acc", test_acc)
mlflow.log_metric("test_loss", test_loss)
# Save as TensorFlow SavedModel flavor
mlflow.keras.log_model(model, "keras-model-tf", save_format="tf")
# Save as default H5 format
mlflow.keras.log_model(model, "keras-model-h5")
# Save as TensorFlow SavedModel format - non-flavor artifact
path = "keras-model-tf-non-flavor"
tf.keras.models.save_model(model, path, overwrite=True, include_optimizer=True)
mlflow.log_artifact(path)
# write model summary
summary = []
model.summary(print_fn=summary.append)
summary = "\n".join(summary)
with open("model_summary.txt", "w") as f:
f.write(summary)
mlflow.log_artifact("model_summary.txt")
else:
# utils.register_model(run, model_name)
pass
# write model as yaml file
with open("model.yaml", "w") as f:
f.write(model.to_yaml())
mlflow.log_artifact("model.yaml")
# MLflow - log onnx model
if log_as_onnx:
import onnx_utils
mname = f"{model_name}_onnx" if model_name else None
onnx_utils.log_model(model, "onnx-model", mname)
# predictions = model.predict_classes(x_test)
predictions = np.argmax(model.predict(x_test), axis=-1)
print("predictions:", predictions)
def train(data, max_depth, max_bins, model_name, log_as_mleap, log_as_onnx):
(trainingData, testData) = data.randomSplit([0.7, 0.3], 42)
print("testData.schema:")
testData.printSchema()
# MLflow - log parameters
print("Parameters:")
print(" max_depth:", max_depth)
print(" max_bins:", max_bins)
mlflow.log_param("max_depth", max_depth)
mlflow.log_param("max_bins", max_bins)
# Create pipeline
dt = DecisionTreeRegressor(labelCol=colLabel,
featuresCol=colFeatures,
maxDepth=max_depth,
maxBins=max_bins)
assembler = VectorAssembler(inputCols=data.columns[:-1],
outputCol=colFeatures)
pipeline = Pipeline(stages=[assembler, dt])
# Fit model and predict
model = pipeline.fit(trainingData)
predictions = model.transform(testData)
# MLflow - log metrics
print("Metrics:")
predictions = model.transform(testData)
metrics = ["rmse", "r2", "mae"]
for metric_name in metrics:
evaluator = RegressionEvaluator(labelCol=colLabel,
predictionCol=colPrediction,
metricName=metric_name)
metric_value = evaluator.evaluate(predictions)
print(f" {metric_name}: {metric_value}")
mlflow.log_metric(metric_name, metric_value)
# MLflow - log spark model
mlflow.spark.log_model(model, "spark-model", \
registered_model_name=None if not model_name else f"{model_name}")
# MLflow - log as MLeap model
if log_as_mleap:
scoreData = testData.drop("quality")
mlflow.mleap.log_model(spark_model=model, sample_input=scoreData, artifact_path="mleap-model", \
registered_model_name=None if not model_name else f"{model_name}_mleap")
# Log MLeap schema file for MLeap runtime deserialization
schema_path = "schema.json"
with open(schema_path, 'w') as f:
f.write(scoreData.schema.json())
print("schema_path:", schema_path)
mlflow.log_artifact(schema_path, "mleap-model")
# MLflow - log as ONNX model
if log_as_onnx:
import onnx_utils
scoreData = testData.drop("quality")
onnx_utils.log_model(spark, model, "onnx-model", model_name, scoreData)
def train(run, model_name, data_path, epochs, batch_size, mlflow_custom_log,
log_as_onnx):
x_train, y_train, x_test, y_test = utils.get_train_data(data_path)
model = build_model()
model.compile(optimizer="rmsprop",
loss="categorical_crossentropy",
metrics=["accuracy"])
model.summary()
model.fit(x_train,
y_train,
epochs=epochs,
batch_size=batch_size,
verbose=0)
print("model.type:", type(model))
test_loss, test_acc = model.evaluate(x_test, y_test)
print("test_acc:", test_acc)
print("test_loss:", test_loss)
if mlflow_custom_log:
mlflow.log_param("epochs", epochs)
mlflow.log_param("batch_size", batch_size)
mlflow.log_metric("test_acc", test_acc)
mlflow.log_metric("test_loss", test_loss)
# Save as TensorFlow SavedModel format (MLflow Keras default)
mlflow.keras.log_model(model,
"keras-model",
registered_model_name=model_name)
#mlflow.keras.log_model(model, "keras-model")
# write model summary
summary = []
model.summary(print_fn=summary.append)
summary = "\n".join(summary)
with open("model_summary.txt", "w") as f:
f.write(summary)
mlflow.log_artifact("model_summary.txt")
elif model_name:
utils.register_model(run, model_name)
# write model as yaml file
with open("model.yaml", "w") as f:
f.write(model.to_yaml())
mlflow.log_artifact("model.yaml")
# MLflow - log onnx model
if log_as_onnx:
import onnx_utils
mname = f"{model_name}_onnx" if model_name else None
onnx_utils.log_model(model, "onnx-model", mname)
predictions = model.predict_classes(x_test)
print("predictions:", predictions)
def train(epochs, batch_size, autolog, log_as_onnx):
print("autolog:", autolog)
x_train, y_train, x_test, y_test = utils.build_data()
model = build_model()
print("model:", type(model))
model.compile(optimizer='rmsprop',
loss='categorical_crossentropy',
metrics=['accuracy'])
model.summary()
model.fit(x_train,
y_train,
epochs=epochs,
batch_size=batch_size,
verbose=0)
test_loss, test_acc = model.evaluate(x_test, y_test)
print("test_acc:", test_acc)
print("test_loss:", test_loss)
if not autolog:
mlflow.log_param("epochs", epochs)
mlflow.log_param("batch_size", batch_size)
mlflow.log_metric("test_acc", test_acc)
mlflow.log_metric("test_loss", test_loss)
mlflow.keras.log_model(model, "keras-model")
#mlflow.tensorflow.log_model(model, "tensorflow-model")
# write model as yaml file
with open("model.yaml", "w") as f:
f.write(model.to_yaml())
mlflow.log_artifact("model.yaml")
# write model summary
summary = []
model.summary(print_fn=summary.append)
summary = '\n'.join(summary)
with open("model_summary.txt", "w") as f:
f.write(summary)
mlflow.log_artifact("model_summary.txt")
# MLflow - log onnx model
if log_as_onnx:
import onnx_utils
onnx_utils.log_model(model, "onnx-model")
predictions = model.predict_classes(x_test)
print("predictions:", predictions)
def train(run, model_name, epochs, batch_size, mlflow_custom_log, log_as_onnx):
x_train, y_train, x_test, y_test = utils.build_data()
model = build_model()
model.compile(
optimizer='rmsprop',
loss='categorical_crossentropy',
metrics=['accuracy'])
model.summary()
model.fit(x_train, y_train, epochs=epochs, batch_size=batch_size, verbose=0)
test_loss, test_acc = model.evaluate(x_test, y_test)
print("test_acc:", test_acc)
print("test_loss:", test_loss)
if mlflow_custom_log:
mlflow.log_param("epochs", epochs)
mlflow.log_param("batch_size", batch_size)
mlflow.log_metric("test_acc", test_acc)
mlflow.log_metric("test_loss", test_loss)
mlflow.keras.log_model(model, "keras-model", registered_model_name=model_name)
# write model summary
summary = []
model.summary(print_fn=summary.append)
summary = '\n'.join(summary)
with open("model_summary.txt", "w") as f:
f.write(summary)
mlflow.log_artifact("model_summary.txt")
else:
utils.register_model(run, model_name)
# write model as yaml file
with open("model.yaml", "w") as f:
f.write(model.to_yaml())
mlflow.log_artifact("model.yaml")
# MLflow - log onnx model
if log_as_onnx:
import onnx_utils
mname = f"{model_name}_onnx" if model_name else None
onnx_utils.log_model(model, "onnx-model", mname)
predictions = model.predict_classes(x_test)
print("predictions:", predictions)
def train(run, model_name, data_path, epochs, batch_size, mlflow_custom_log,
log_as_onnx, log_as_tensorflow_lite, log_as_tensorflow_js):
print("mlflow_custom_log:", mlflow_custom_log)
x_train, _, y_train, _ = utils.build_data(data_path)
ncols = x_train.shape[1]
def baseline_model():
model = Sequential()
model.add(
Dense(ncols,
input_dim=ncols,
kernel_initializer='normal',
activation='relu'))
model.add(Dense(1, kernel_initializer='normal'))
model.compile(loss='mean_squared_error', optimizer='adam')
return model
model = baseline_model()
if mlflow_custom_log:
print("Logging with mlflow.log")
mlflow.log_param("epochs", epochs)
mlflow.log_param("batch_size", batch_size)
mlflow.keras.log_model(model,
"tensorflow-model",
registered_model_name=model_name)
else:
utils.register_model(run, model_name)
# MLflow - log as ONNX model
if log_as_onnx:
import onnx_utils
mname = f"{model_name}_onnx" if model_name else None
onnx_utils.log_model(model, "onnx-model", model_name=mname)
# Save as TensorFlow Lite format
if log_as_tensorflow_lite:
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
path = "model.tflite"
with open(path, "wb") as f:
f.write(tflite_model)
mlflow.log_artifact(path, "tensorflow-lite-model")
# Save as TensorFlow.js format
if log_as_tensorflow_js:
import tensorflowjs as tfjs
path = "model.tfjs"
tfjs.converters.save_keras_model(model, path)
mlflow.log_artifact(path, "tensorflow-js-model")
# Evaluate model
estimator = KerasRegressor(build_fn=baseline_model,
epochs=epochs,
batch_size=batch_size,
verbose=0)
kfold = KFold(n_splits=10)
results = cross_val_score(estimator, x_train, y_train, cv=kfold)
print(
f"Baseline MSE: mean: {round(results.mean(),2)} std: {round(results.std(),2)}"
)
if mlflow_custom_log:
mlflow.log_metric("mse_mean", results.mean())
mlflow.log_metric("mse_std", results.std())
# Score
data = x_train
predictions = model.predict(data)
predictions = pd.DataFrame(data=predictions, columns=["prediction"])
print("predictions.shape:", predictions.shape)
print("predictions:", predictions)