如何从40,40,9要素和40,40目标数据集制作训练和测试数据集

Question

我想训练数据集并测试数据集
特征形状为（40,40,9），目标形状为（40,40）

def trainGenerator():

    train_path = '/media/jake/mark-4tb3/input/dacon_4tb/ai_friends_rain/train'
    train_files = sorted(glob.glob(train_path + '/*'))

    for file in train_files:

        dataset = np.load(file)

        target= dataset[:,:,-1].reshape(40,40,1)
        cutoff_labels = np.where(target < 0, 0, target)
        feature = dataset[:,:,:9]

        if (cutoff_labels > 0).sum() < 10:

            continue

        yield (feature, cutoff_labels)

train_dataset = tf.data.Dataset.from_generator(trainGenerator, (tf.float32, tf.float32), (tf.TensorShape([40,40,9]),tf.TensorShape([40,40,1])))

这是使用tensorflow制作数据集的一种方法
但是我想分别坐火车和测试仪来训练Unet

dataset = np.load(train_files[0])
print(dataset.shape)
print(dataset[0].shape)
target = dataset[:,:,-1]
print('target->',target.shape)
feature = dataset[:,:,:9]
print('feature->',feature.shape)
(40, 40, 15)
(40, 15)
target-> (40, 40)
feature-> (40, 40, 9)
time: 3.92 ms

不确定火车和目标的形状会是吗？
如果有10000张图像
火车：（40,40,9,10000）和目标（40,40,10000）？