40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
import mlflow
|
|
import mlflow.keras
|
|
import pandas as pd
|
|
import numpy as np
|
|
from tensorflow.keras import Sequential
|
|
from tensorflow.keras.layers import Dense
|
|
from sklearn.preprocessing import MinMaxScaler
|
|
import sys
|
|
|
|
# Parameters from the command line
|
|
epochs = int(sys.argv[1])
|
|
batch_size = int(sys.argv[2])
|
|
|
|
mlflow.start_run()
|
|
|
|
train_data = pd.read_csv('./data/car_prices_train.csv')
|
|
train_data.dropna(inplace=True)
|
|
y_train = train_data['sellingprice'].astype(np.float32)
|
|
X_train = train_data[['year', 'condition', 'transmission']]
|
|
|
|
scaler_x = MinMaxScaler()
|
|
X_train['condition'] = scaler_x.fit_transform(X_train[['condition']])
|
|
|
|
scaler_y = MinMaxScaler()
|
|
y_train = scaler_y.fit_transform(y_train.values.reshape(-1, 1))
|
|
X_train = pd.get_dummies(X_train, columns=['transmission'])
|
|
|
|
model = Sequential([Dense(64, activation='relu'), Dense(32, activation='relu'), Dense(1)])
|
|
model.compile(optimizer='adam', loss='mean_squared_error')
|
|
|
|
# Training the model with MLflow tracking
|
|
model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size)
|
|
|
|
mlflow.keras.log_model(model, "model")
|
|
|
|
mlflow.log_param("epochs", epochs)
|
|
mlflow.log_param("batch_size", batch_size)
|
|
|
|
mlflow.end_run()
|