用于 CNN 的 ImageDataGenerator()，输入和输出为图像

Question

I'm looking for a training map with something like this:

Grayscale Image -> Coloured Image

But the dataset can't be loaded all to the ram as X and Y because of obvious reasons.

I looked up the ImageDataGenerator() library, but it didn't give me a clear answer as to make it work here.

Summary:

Input Shape = (2048, 2048, 1)

Output Shape = (2048, 2048, 2)

Training Dataset = 17,000 images

Validation Dataset = 1,000 images

Here's the structure of the model I'm trying to train:

Model: "functional_1"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 2048, 2048,  0                                            
__________________________________________________________________________________________________
conv2d (Conv2D)                 (None, 2048, 2048, 1 160         input_1[0][0]                    
__________________________________________________________________________________________________
leaky_re_lu (LeakyReLU)         (None, 2048, 2048, 1 0           conv2d[0][0]                     
__________________________________________________________________________________________________
conv2d_1 (Conv2D)               (None, 2048, 2048, 3 4640        leaky_re_lu[0][0]                
__________________________________________________________________________________________________
leaky_re_lu_1 (LeakyReLU)       (None, 2048, 2048, 3 0           conv2d_1[0][0]                   
__________________________________________________________________________________________________
batch_normalization (BatchNorma (None, 2048, 2048, 3 128         leaky_re_lu_1[0][0]              
__________________________________________________________________________________________________
max_pooling2d (MaxPooling2D)    (None, 1024, 1024, 3 0           batch_normalization[0][0]        
__________________________________________________________________________________________________
conv2d_2 (Conv2D)               (None, 1024, 1024, 6 18496       max_pooling2d[0][0]              
__________________________________________________________________________________________________
leaky_re_lu_2 (LeakyReLU)       (None, 1024, 1024, 6 0           conv2d_2[0][0]                   
__________________________________________________________________________________________________
batch_normalization_1 (BatchNor (None, 1024, 1024, 6 256         leaky_re_lu_2[0][0]              
__________________________________________________________________________________________________
max_pooling2d_1 (MaxPooling2D)  (None, 512, 512, 64) 0           batch_normalization_1[0][0]      
__________________________________________________________________________________________________
conv2d_3 (Conv2D)               (None, 512, 512, 128 73856       max_pooling2d_1[0][0]            
__________________________________________________________________________________________________
leaky_re_lu_3 (LeakyReLU)       (None, 512, 512, 128 0           conv2d_3[0][0]                   
__________________________________________________________________________________________________
batch_normalization_2 (BatchNor (None, 512, 512, 128 512         leaky_re_lu_3[0][0]              
__________________________________________________________________________________________________
conv2d_4 (Conv2D)               (None, 512, 512, 256 295168      batch_normalization_2[0][0]      
__________________________________________________________________________________________________
leaky_re_lu_4 (LeakyReLU)       (None, 512, 512, 256 0           conv2d_4[0][0]                   
__________________________________________________________________________________________________
batch_normalization_3 (BatchNor (None, 512, 512, 256 1024        leaky_re_lu_4[0][0]              
__________________________________________________________________________________________________
up_sampling2d (UpSampling2D)    (None, 1024, 1024, 2 0           batch_normalization_3[0][0]      
__________________________________________________________________________________________________
conv2d_5 (Conv2D)               (None, 1024, 1024, 1 295040      up_sampling2d[0][0]              
__________________________________________________________________________________________________
leaky_re_lu_5 (LeakyReLU)       (None, 1024, 1024, 1 0           conv2d_5[0][0]                   
__________________________________________________________________________________________________
batch_normalization_4 (BatchNor (None, 1024, 1024, 1 512         leaky_re_lu_5[0][0]              
__________________________________________________________________________________________________
up_sampling2d_1 (UpSampling2D)  (None, 2048, 2048, 1 0           batch_normalization_4[0][0]      
__________________________________________________________________________________________________
conv2d_6 (Conv2D)               (None, 2048, 2048, 6 73792       up_sampling2d_1[0][0]            
__________________________________________________________________________________________________
leaky_re_lu_6 (LeakyReLU)       (None, 2048, 2048, 6 0           conv2d_6[0][0]                   
__________________________________________________________________________________________________
concatenate (Concatenate)       (None, 2048, 2048, 6 0           leaky_re_lu_6[0][0]              
                                                                 input_1[0][0]                    
__________________________________________________________________________________________________
conv2d_7 (Conv2D)               (None, 2048, 2048, 6 37504       concatenate[0][0]                
__________________________________________________________________________________________________
leaky_re_lu_7 (LeakyReLU)       (None, 2048, 2048, 6 0           conv2d_7[0][0]                   
__________________________________________________________________________________________________
batch_normalization_5 (BatchNor (None, 2048, 2048, 6 256         leaky_re_lu_7[0][0]              
__________________________________________________________________________________________________
conv2d_8 (Conv2D)               (None, 2048, 2048, 3 18464       batch_normalization_5[0][0]      
__________________________________________________________________________________________________
leaky_re_lu_8 (LeakyReLU)       (None, 2048, 2048, 3 0           conv2d_8[0][0]                   
__________________________________________________________________________________________________
conv2d_9 (Conv2D)               (None, 2048, 2048, 2 578         leaky_re_lu_8[0][0]              
==================================================================================================
Total params: 820,386
Trainable params: 819,042
Non-trainable params: 1,344
__________________________________________________________________________________________________

Answer 1

I am gonna post my comment with an example code:

You need to train the model with batches. If for example, you want to use 500 images in 1 epoch, you can do 50 images per batches and 10 epochs instead. That way you are only loading 50 images in memory. You have to set the batch size and set shuffle to True so your batches have different images.

If you have your images in a directory, the comment above translate into something like this:

from keras_preprocessing.image import ImageDataGenerator
preprocessing_images = ImageDataGenerator()

train_generator = preprocessing_images.flow_from_directory(
        train_path,
        target_size=target_size,
        batch_size=50,
        class_mode="categorical", # classes are provided in categorical format for a 2-unit output layer
        shuffle=True,
        color_mode="grayscale",
        seed=1234567890
    )

There you have a generator which will yield batches of 50 images. I have not specified the arguments for your specific problem, you need to change the target_size and all that stuff. Note that this is an infinite generator, it is going to return an infinite number of batches with 50 images. You can wrap it into a tf.data.Dataset or specify the steps per epoch. There are several ways. I hope it helps. Work a bit with it and if you still have problems, I will elaborate on the answer(I am a busy right now)

Answer 2

That would be easiest with a custom training loop.

def reconstruct(colored_inputs):
    with tf.GradientTape() as tape:
        grayscale_inputs = tf.image.rgb_to_grayscale(colored_inputs)

        out = autoencoder(grayscale_inputs)
        loss = loss_object(colored_inputs, out)

    gradients = tape.gradient(loss, autoencoder.trainable_variables)
    optimizer.apply_gradients(zip(gradients, autoencoder.trainable_variables))

    reconstruction_loss(loss)

Here, my data iterator is cyling through all the color pictures, but its converted to grayscale before being passed to the model. Then, the RGB output of the model is compared to the original RGB image. You will have to use the argument class_mode=None in flow_from_directory . I used tf.image.rgb_to_grayscale to make the conversion between grayscale and RGB.

Full example:

import tensorflow as tf
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
import os

os.chdir(r'catsanddogs')

generator = tf.keras.preprocessing.image.ImageDataGenerator()
iterator = generator.flow_from_directory(
    target_size=(32, 32),
    directory='.',
    batch_size=4,
    class_mode=None)

encoder = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(32, 32, 1)),
    tf.keras.layers.Dense(32),
    tf.keras.layers.Dense(16)
])

decoder = tf.keras.Sequential([
    tf.keras.layers.Dense(32, input_shape=[16]),
    tf.keras.layers.Dense(32 * 32 * 3),
    tf.keras.layers.Reshape([32, 32, 3])
])


autoencoder = tf.keras.Sequential([encoder, decoder])

loss_object = tf.losses.BinaryCrossentropy()

reconstruction_loss = tf.metrics.Mean(name='reconstruction_loss')

optimizer = tf.optimizers.Adam()


def reconstruct(colored_inputs):
    with tf.GradientTape() as tape:
        grayscale_inputs = tf.image.rgb_to_grayscale(colored_inputs)

        out = autoencoder(grayscale_inputs)
        loss = loss_object(colored_inputs, out)

    gradients = tape.gradient(loss, autoencoder.trainable_variables)
    optimizer.apply_gradients(zip(gradients, autoencoder.trainable_variables))

    reconstruction_loss(loss)


if __name__ == '__main__':
    template = 'Epoch {:2} Reconstruction Loss {:.4f}'
    for epoch in range(50):
        reconstruction_loss.reset_states()
        for input_batches in iterator:
            reconstruct(input_batches)
        print(template.format(epoch + 1, reconstruction_loss.result()))