Ускорить обучение

from tensorflow.keras import layers
from tensorflow.keras import Input
from tensorflow.keras.models import Model
from tensorflow.keras import optimizers
import numpy as np
import os
import glob
import matplotlib.pyplot as plt
%matplotlib inline
 
latent_dim = 128
height = 128
width = 128
channels = 3
 
generator_input = Input(shape=(latent_dim,))
 
x = layers.Dense(512 * 64 * 64)(generator_input)
x = layers.LeakyReLU()(x)
x = layers.Reshape((64, 64, 512))(x)
x = layers.Conv2D(1024, 5, padding='same')(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2DTranspose(1024, 4, strides=2, padding='same')(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(1024, 5, padding='same')(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(1024, 5, padding='same')(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(channels, 7, activation='tanh', padding='same')(x)
generator = Model(generator_input, x)
 
discriminator_input = layers.Input(shape=(height, width, channels))
x = layers.Conv2D(512, 3)(discriminator_input)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(512, 4, strides=2)(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(512, 4, strides=2)(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(512, 4, strides=2)(x)
x = layers.LeakyReLU()(x)
x = layers.Flatten()(x)
x = layers.Dropout(0.4)(x)
x = layers.Dense(1, activation='sigmoid')(x)
 
discriminator = Model(discriminator_input, x)
 
discriminator_optimizer = optimizers.legacy.RMSprop(learning_rate=0.0008, clipvalue=1.0, decay=1e-8)
discriminator.compile(optimizer=discriminator_optimizer, loss='binary_crossentropy')
 
discriminator.trainable = False
gan_input = Input(shape=(latent_dim,))
gan_output = discriminator(generator(gan_input))
gan = Model(gan_input, gan_output)
gan_optimizer = optimizers.legacy.RMSprop(learning_rate=0.0004, clipvalue=1.0, decay=1e-8)
gan.compile(optimizer=gan_optimizer, loss='binary_crossentropy')
 
DATA_PATH = r'C:\Users\Acer\Desktop\fragments'
extensions = [DATA_PATH + '/*.png']
listOfFiles = []
 
for i in extensions:
    for image_file_name in glob.glob(i):    
        listOfFiles.append(image_file_name)
 
x_train = np.zeros(shape = (len(listOfFiles), height, width, channels), dtype = 'float32')
 
for id, file in enumerate(listOfFiles):
    x_train[id] = plt.imread(file)[:,:,:3]
 
iterations = 10000
batch_size = 2
 
start = 0
for step in range(iterations):
    random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))
    
    #Можно чуть чуть ускорить на этой строке
    generated_images = generator.predict(random_latent_vectors)
    
    stop = start + batch_size
    real_images = x_train[start: stop]    
    combined_images = np.concatenate([generated_images, real_images])    
    labels = np.concatenate([np.ones((batch_size, 1)), np.zeros((batch_size, 1))])    
    labels += 0.05 * np.random.random(labels.shape)    
    d_loss = discriminator.train_on_batch(combined_images, labels)    
    random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))    
    misleading_targets = np.zeros((batch_size, 1))
    
    #На этой строке надо ускорить сильно, почти всю RAM (16 гига) использует
    a_loss = gan.train_on_batch(random_latent_vectors, misleading_targets)
    
    start += batch_size
    if start > len(x_train) - batch_size:
        start = 0