jordon/DeepEncode
Archived
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Added GPU growth

Browse source
This commit is contained in:
Jordon Brooks 2023-08-17 01:57:43 +01:00
parent 046b78cb87
commit 3ea1568ad3
1 changed files with 11 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

12
train_model.py
View file

@ -14,8 +14,18 @@ from featureExtraction import psnr
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
import gc
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, Callback
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
except RuntimeError as e:
print(e)
from video_compression_model import VideoCompressionModel, data_generator