48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
import pandas as pd
|
|
from os import path
|
|
from tensorflow import keras
|
|
from tensorflow.keras import layers
|
|
import mlflow
|
|
import sys
|
|
|
|
model_name = "model.h5"
|
|
train_data=pd.read_csv('train.csv')
|
|
|
|
input_columns=["Age","Nationality","Position","Club"]
|
|
X=train_data[input_columns].to_numpy()
|
|
Y=train_data[["Overall"]].to_numpy()
|
|
model = None
|
|
|
|
model = keras.Sequential(name="fifa_overall")
|
|
model.add(keras.Input(shape=(len(input_columns),), name="player_info"))
|
|
model.add(layers.Dense(4, activation="relu", name="layer1"))
|
|
model.add(layers.Dense(8, activation="relu", name="layer2"))
|
|
model.add(layers.Dense(8, activation="relu", name="layer3"))
|
|
model.add(layers.Dense(5, activation="relu", name="layer4"))
|
|
model.add(layers.Dense(1, activation="relu", name="output"))
|
|
|
|
model.compile(
|
|
optimizer=keras.optimizers.RMSprop(),
|
|
loss=keras.losses.MeanSquaredError(),
|
|
)
|
|
|
|
batch_size = int(sys.argv[1])
|
|
epochs = int(sys.argv[2])
|
|
mlflow.log_param("batch_size", batch_size)
|
|
mlflow.log_param("epochs", epochs)
|
|
|
|
history = model.fit(
|
|
X,
|
|
Y,
|
|
batch_size=batch_size,
|
|
epochs=epochs,
|
|
)
|
|
model.save(model_name)
|
|
signature = mlflow.models.signature.infer_signature(X, model.predict(X))
|
|
input_example = {
|
|
"Age": 0.38,
|
|
"Nationality": 0.175,
|
|
"Position": 1,
|
|
"Club" :0.91846154
|
|
}
|
|
mlflow.keras.save_model(model, "fifa_overall", signature=signature, input_example = input_example) |