55 lines
2.1 KiB
Python
55 lines
2.1 KiB
Python
|
import sys
|
||
|
|
import pandas as pd
|
||
|
|
from tensorflow import keras
|
||
|
|
from sklearn.model_selection import train_test_split
|
||
|
|
from sklearn.metrics import mean_squared_error as rmse
|
||
|
|
import mlflow as mlf
|
||
|
|
|
||
|
|
|
||
|
|
def create_model(test_size, epochs, batch_size):
|
||
|
|
df = pd.read_csv('country_vaccinations.csv').dropna()
|
||
|
|
dataset = df.iloc[:, 3:-3]
|
||
|
|
dataset = df.groupby(by=["country"], dropna=True).sum()
|
||
|
|
X = dataset.loc[:,dataset.columns != "daily_vaccinations"]
|
||
|
|
y = dataset.loc[:,dataset.columns == "daily_vaccinations"]
|
||
|
|
|
||
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = test_size, random_state = 6)
|
||
|
|
|
||
|
|
model = keras.Sequential([
|
||
|
|
keras.layers.Dense(512,input_dim = X_train.shape[1],kernel_initializer='normal', activation='relu'),
|
||
|
|
keras.layers.Dense(512,kernel_initializer='normal', activation='relu'),
|
||
|
|
keras.layers.Dense(256,kernel_initializer='normal', activation='relu'),
|
||
|
|
keras.layers.Dense(256,kernel_initializer='normal', activation='relu'),
|
||
|
|
keras.layers.Dense(128,kernel_initializer='normal', activation='relu'),
|
||
|
|
keras.layers.Dense(1,kernel_initializer='normal', activation='linear'),
|
||
|
|
])
|
||
|
|
|
||
|
|
model.compile(loss='mean_absolute_error', optimizer='adam', metrics=['mean_absolute_error'])
|
||
|
|
|
||
|
|
model.fit(X_train, y_train, epochs = epochs, validation_split = 0.3, batch_size = batch_size)
|
||
|
|
|
||
|
|
prediction = model.predict(X_test)
|
||
|
|
rmse_result = rmse(y_test, prediction, squared = False)
|
||
|
|
print(prediction)
|
||
|
|
model.save('vaccines_model')
|
||
|
|
return model, rmse_result
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
test_size = float(sys.argv[1]) if len(sys.argv) > 1 else 0.2
|
||
|
|
epochs = int(sys.argv[2]) if len(sys.argv) > 1 else 100
|
||
|
|
batch_size = int(sys.argv[3]) if len(sys.argv) > 1 else 32
|
||
|
|
|
||
|
|
with mlf.start_run():
|
||
|
|
mlf.log_param("Test size", test_size)
|
||
|
|
mlf.log_param("Epochs", epochs)
|
||
|
|
mlf.log_param("Batch size", batch_size)
|
||
|
|
|
||
|
|
model, rmse_result = create_model(
|
||
|
|
test_size=test_size,
|
||
|
|
epochs=epochs,
|
||
|
|
batch_size=batch_size,
|
||
|
|
)
|
||
|
|
|
||
|
|
mlf.log_metric("RMSE", rmse_result)
|
||
|
|
|
||
|
|
mlf.keras.log_model(model, "country_vaccination")
|