ium_434704/training.py
Wojciech Jarmosz f1083d4e7a
Some checks failed
s434704-training/pipeline/head There was a failure building this commit
Add evaluation
2021-05-14 04:01:11 +02:00

51 lines
1.6 KiB
Python

import sys
import pandas as pd
import numpy as np
import tensorflow as tf
import os.path
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.layers.experimental import preprocessing
arguments = sys.argv[1:]
verbose = [command.split('=')[1] for command in arguments if command.split('=')[0] == 'verbose']
epochs = [command.split('=')[1] for command in arguments if command.split('=')[0] == 'epochs']
pd.set_option("display.max_columns", None)
# Wczytanie danych
train_data = pd.read_csv("./MoviesOnStreamingPlatforms_updated.train")
# Stworzenie modelu
columns_to_use = ['Year', 'Runtime', 'Netflix']
train_X = tf.convert_to_tensor(train_data[columns_to_use])
train_Y = tf.convert_to_tensor(train_data[["IMDb"]])
normalizer = preprocessing.Normalization(input_shape=[3,])
normalizer.adapt(train_X)
if os.path.isdir('linear_regression'):
model = keras.models.load_model('linear_regression')
else:
model = keras.Sequential([
keras.Input(shape=(len(columns_to_use),)),
normalizer,
layers.Dense(30, activation='relu'),
layers.Dense(10, activation='relu'),
layers.Dense(25, activation='relu'),
layers.Dense(1)
])
model.compile(loss='mean_absolute_error',
optimizer=tf.keras.optimizers.Adam(0.001),
metrics=[tf.keras.metrics.RootMeanSquaredError()])
verbose = 0 if len(verbose) == 0 else verbose[0]
epochs = 100 if len(epochs) == 0 else epochs[0]
model.fit(train_X, train_Y, verbose=verbose, epochs=epochs)
model.save('linear_regression.h5')