ium_434704/training.py

45 lines
1.2 KiB
Python
Raw Permalink Normal View History

import sys
2021-05-13 22:20:25 +02:00
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:]
2021-05-15 09:06:59 +02:00
verbose = int(arguments[0])
epochs = int(arguments[1])
2021-05-13 22:20:25 +02:00
pd.set_option("display.max_columns", None)
# Wczytanie danych
2021-05-14 00:58:33 +02:00
train_data = pd.read_csv("./MoviesOnStreamingPlatforms_updated.train")
2021-05-13 22:20:25 +02:00
# 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)
2021-05-15 03:50:19 +02:00
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()])
model.fit(train_X, train_Y, verbose=verbose, epochs=epochs)
model.save('linear_regression.h5')