python - 具有多个输入的 TensorFlow 生成器读取多个文件会导致内核崩溃

Question

我是 StackOverflow 以及生成器的新手，所以这里可能有一些糟糕的格式选择。我正在做一个机器学习项目，我在阅读中发现定义一个生成器对我的数据处理很有帮助，因为我的数据集非常大，我需要做很多预处理。我的模型有两个输入，每个输入一个形状矩阵 (5,480,480,3)，并输出 9 个类中的一个。我的数据集是一堆 npz 文件，其中包含 20 (5,480,480,3) 个矩阵。我的生成器需要调用一个名为get_sub_images（我已经充分调试并确定不是问题）将其中一个矩阵转换为相同类型的矩阵（此函数使用 YOLO 识别具有边界框的对象，并将 5 个图像中的每一个拆分为对象和背景图像）。如果我将批处理大小设为一个文件，此生成器工作正常，但是当我将批处理大小增加到 2 时内核崩溃，我不知道为什么。这可能是什么原因造成的？我认为这要么与我的循环结构中的错误有关，要么与我的预处理功能所花费的时间有关（如果这可能会导致生成器出现问题）。

我在下面附上了我的相关代码。首先是我的生成器功能。

def tf_data_generator(file_list, batch_size = 1):
i = 0
while True:
    if i*batch_size >= len(file_list):
        print("\nShuffling and Restarting...")
        i = 0
        np.random.shuffle(file_list)
    else:
        file_chunk = file_list[i*batch_size:(i+1)*batch_size] 
        print("\nTraining Files: ", file_chunk)
        data = [[],[]]
        labels = []
        
        for fil in file_chunk:
            with np.load(fil) as dat:
                training_data_images = dat['data']
                training_data_labels = dat['labels']


            to_shuffle = list(zip(training_data_images, training_data_labels))
            np.random.shuffle(to_shuffle)
            training_data_images, training_data_labels = zip(*to_shuffle)

            train_objects = []
            train_background = []
            for k in range(len(training_data_images)):
                train_objects.append([])
                train_background.append([])
                for j in range(5):
                    temp1,temp2 = get_sub_images(training_data_images[k][j])
                    train_objects[k].append(temp1)
                    train_background[k].append(temp2)

            X_train = [reshape_custom_X(np.array(train_objects)), reshape_custom_X(np.array(train_background))]
            y_train = reshape_custom_y(np.array(training_data_labels))
            
            for x in range(X_train[0].shape[0]):
                data[0].append(X_train[0][x])
                data[1].append(X_train[1][x])
                labels.append(y_train[x])
            
            
        yield (np.array(data[0]), np.array(data[1])), labels
        i = i + 1

这是我的培训电话。更改batch_size = 2会导致我的内核崩溃。

epochs = 1
batch_size = 1
file_names = glob.glob("Data\\**\*.npz", recursive=True)
np.random.shuffle(file_names)
train_file_names = file_names[0:int(len(file_names)*0.7)]
validation_file_names = file_names[int(len(file_names)*0.7):]

print(len(train_file_names), len(validation_file_names))
print(train_file_names, "\n", validation_file_names)


train_dataset = tf.data.Dataset.from_generator(tf_data_generator, args = [train_file_names, batch_size], 
                                              output_shapes = (((None,5,416,416,3), (None,5,416,416,3)),(None,9)),
                                              output_types = ((tf.float32, tf.float32), tf.float32))

validation_dataset = tf.data.Dataset.from_generator(tf_data_generator, args = [validation_file_names, batch_size],
                                              output_shapes = (((None,5,416,416,3), (None,5,416,416,3)),(None,9)),
                                              output_types = ((tf.float32, tf.float32), tf.float32))


steps_per_epoch = np.int(np.ceil(len(train_file_names)/batch_size))
validation_steps = np.int(np.ceil(len(validation_file_names)/batch_size))
print("steps_per_epoch = ", steps_per_epoch)
print("validation_steps = ", validation_steps)

model.fit(train_dataset, validation_data = validation_dataset, steps_per_epoch = steps_per_epoch,
         validation_steps = validation_steps, epochs = epochs)

最后，这里有一些关于我的模型的相关信息，以防万一。我已经包含了我的编译语句以及我的第一层和最后一层。

input1 = tf.keras.Input(shape=(5, 416, 416, 3))
input2 = tf.keras.Input(shape=(5, 416, 416, 3))
out = Dense(9, activation='softmax')(out)

model = Model(inputs=[x1.input,x2.input], outputs=out)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

谢谢您的帮助。

标签： pythontensorflowgenerator