Ошибка TensorFlow из-за нехватки памяти при работе с графическим процессором ⇐ Python

[Anonymous]

Я скопировал простой пример автокодировщика из Интернета и установил Tensorflow 2.10.1 на Windows WSL2 с помощью этого руководства: Установите TensorFlow с помощью pip — WSL2 (tensorflow.org)
Мой вывод для nvidia-smi:


Это мой автоэнкодер:

Код: Выделить всё

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"

from tensorflow.keras.models import load_model # type: ignore
import random
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow import keras
K = keras.backend

from tensorflow.keras import mixed_precision

# Enable mixed precision
mixed_precision.set_global_policy('mixed_float16')

BATCH_SIZE = 8

IMG_WIDTH = 96
IMG_HEIGHT = 96

ENCODING_SIZE = 8

tf.device('/gpu:0')

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# Set memory growth to true
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
except RuntimeError as e:
print(e)

flower_generator = tf.keras.preprocessing.image_dataset_from_directory(
"/mnt/c/dev/autoencoder/input/flowers",
seed=84,
image_size=(IMG_HEIGHT, IMG_WIDTH),
batch_size=BATCH_SIZE,
labels='inferred'
)

# keep the flower names for future reference
flower_names = flower_generator.class_names

# create a training data generator
training_generator = tf.keras.preprocessing.image_dataset_from_directory(
"/mnt/c/dev/autoencoder/input/flowers",
seed=42,
image_size=(IMG_HEIGHT, IMG_WIDTH),
batch_size=BATCH_SIZE,
labels=None,
validation_split=0.2,
subset='training'
)

# and a validation set generator
validation_generator = tf.keras.preprocessing.image_dataset_from_directory(
"/mnt/c/dev/autoencoder/input/flowers",
seed=42,
image_size=(IMG_HEIGHT, IMG_WIDTH),
batch_size=BATCH_SIZE,
labels=None,
validation_split=0.2,
subset='validation'
)

# autoencoders require the input matrix as both the input
# and the output
# provide input ->  input for the autoencoder
# and, do scaling here

def replicate_inputs_to_outputs(images):
return images/255, images/255

# create prefetch generators to speed things up a bit
training_generator = training_generator.prefetch(128)
validation_generator = validation_generator.prefetch(128)

# create the final generators to be used in training the autoencoders
X_train = training_generator.map(replicate_inputs_to_outputs)
X_valid = validation_generator.map(replicate_inputs_to_outputs)

ENCODING_SIZE = 1024

# build a convolutional auto-encoder
encoder = keras.models.Sequential([
keras.layers.Conv2D(32, input_shape=[IMG_WIDTH, IMG_HEIGHT, 3], kernel_size=3, padding='same', activation='relu'),
keras.layers.MaxPool2D(),
keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='relu'),
keras.layers.MaxPool2D(),
keras.layers.Conv2D(128, kernel_size=3, padding='same', activation='relu'),
keras.layers.Flatten(),
keras.layers.Dense(ENCODING_SIZE)
])

encoder.summary()

decoder = keras.models.Sequential([
keras.layers.Dense(128 * 24 * 24, activation="relu", input_shape=[ENCODING_SIZE]),
keras.layers.Reshape([24, 24, 128]),
keras.layers.Conv2DTranspose(filters=64, kernel_size=3, strides=2, padding="same", activation="relu"),
keras.layers.Conv2DTranspose(filters=32, kernel_size=3, strides=2, padding="same", activation="relu"),
keras.layers.Conv2DTranspose(filters=3, kernel_size=3, padding='same', activation='sigmoid')
])

decoder.summary()

ae = keras.models.Sequential([encoder, decoder])

def exponential_decay_fn(epoch):
return 0.001 * 0.1 ** (epoch / 10)

lr_scheduler = keras.callbacks.LearningRateScheduler(exponential_decay_fn)

# add early stopping
early_stopping_cb = keras.callbacks.EarlyStopping(
patience=5, restore_best_weights=True)

# clear the session for a clean run
keras.backend.clear_session()
tf.random.set_seed(42)

# compile and run
ae.compile(loss="mean_squared_error", optimizer=keras.optimizers.Nadam())
history = ae.fit(X_train, validation_data=X_valid,  epochs=25,
callbacks=[lr_scheduler, early_stopping_cb])
# save model for future use
# handy for fast iteration and experiments
ae.save('trained_ae.h5')
# reload model

ae = load_model('trained_ae.h5')
encoder = ae.layers[0]

plt.figure(figsize=(12, 18))

for images, _ in X_valid.take(1):
ae_images = ae.predict(images)
for i in range(10):
ax = plt.subplot(5, 4, (i * 2) + 1)
plt.imshow(images[i])
plt.title('original')

ax = plt.subplot(5, 4, (i * 2) + 2)
plt.imshow(ae_images[i])
plt.title('reconstructed')
plt.axis("off")

plt.show()

Когда я запускаю это, я получаю кучу следующих сообщений:

И, наконец, выдает ошибку распределения памяти:


Я использовал простой код, чтобы проверить доступную память, и он вернул следующее:

Код: Выделить всё

import subprocess as sp
import os

def get_gpu_memory():
command = "nvidia-smi --query-gpu=memory.free --format=csv"
memory_free_info = sp.check_output(command.split()).decode('ascii').split('\n')[:-1][1:]
memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
return memory_free_values

gpu_memory = get_gpu_memory()
print("Free GPU Memory (in MiB):", gpu_memory)

Как Мне не хватает памяти ошибка?

Подробнее здесь: https://stackoverflow.com/questions/790 ... g-with-gpu