File size: 2,219 Bytes
f2d6c5c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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
# Environment settings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Suppress TensorFlow INFO and WARNING messages
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'
# Suppress TensorFlow logging
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()
|