|
import os |
|
import gc |
|
import sys |
|
import numpy as np |
|
import tensorflow as tf |
|
from tensorflow.keras.models import model_from_json |
|
from sklearn.model_selection import train_test_split |
|
import logging |
|
|
|
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' |
|
os.environ['TF_GPU_ALLOCATOR'] = 'cuda_malloc_async' |
|
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' |
|
os.environ['TF_XLA_FLAGS'] = '--tf_xla_auto_jit=2 --tf_xla_enable_xla_devices' |
|
os.environ['XLA_FLAGS'] = '--xla_gpu_cuda_data_dir=/usr/local/cuda --xla_gpu_force_compilation_parallelism=1' |
|
|
|
|
|
logging.getLogger('tensorflow').setLevel(logging.ERROR) |
|
|
|
BATCHSIZE = int(sys.argv[1]) |
|
EPOCHS = int(sys.argv[2]) |
|
fps_file = sys.argv[3] |
|
y_true_file = sys.argv[4] |
|
|
|
def load_model(): |
|
with open('save_model/model_config.json', 'r') as json_file: |
|
model_json = json_file.read() |
|
model = model_from_json(model_json) |
|
model.load_weights('save_model/model_weights.weights.h5') |
|
model.compile(optimizer=tf.keras.optimizers.Adam(), |
|
loss=tf.keras.losses.MeanSquaredError(), |
|
metrics=[tf.keras.metrics.MeanSquaredError(), |
|
tf.keras.metrics.MeanAbsoluteError(), |
|
tf.keras.metrics.RootMeanSquaredError()]) |
|
return model |
|
|
|
def preprocess_data(xtr, ytr): |
|
buffer_size = min(10000, len(xtr)) |
|
dataset = tf.data.Dataset.from_tensor_slices((xtr, ytr)) |
|
dataset = dataset.shuffle(buffer_size=buffer_size).batch(BATCHSIZE).prefetch(tf.data.AUTOTUNE) |
|
return dataset |
|
|
|
def train_model(model, train_dataset, epochs): |
|
model.fit(train_dataset, epochs=epochs, verbose=0) |
|
model.save('save_model/trained_model.keras') |
|
return model |
|
|
|
def clear_gpu_memory(): |
|
tf.keras.backend.clear_session() |
|
gc.collect() |
|
print("GPU memory cleared.", file=sys.stderr) |
|
|
|
if __name__ == "__main__": |
|
fps = np.load(fps_file) |
|
y_true = np.load(y_true_file) |
|
|
|
xtr, xte, ytr, yte = train_test_split(fps, y_true, test_size=0.2, random_state=42) |
|
train_dataset = preprocess_data(xtr, ytr) |
|
|
|
model = load_model() |
|
trained_model = train_model(model, train_dataset, EPOCHS) |
|
|
|
clear_gpu_memory() |
|
|
