如何使用tf.data.Dataset.from_generator()将参数发送到生成器函数?

时间:2018-09-21 11:57:03

标签: python python-3.x tensorflow tensorflow-datasets

我想使用tf.data.Dataset函数创建多个from_generator()。我想向生成器函数(raw_data_gen)发送一个参数。这个想法是生成器函数将根据发送的参数产生不同的数据。这样,我希望raw_data_gen能够提供培训,验证或测试数据。

training_dataset = tf.data.Dataset.from_generator(raw_data_gen, (tf.float32, tf.uint8), ([None, 1], [None]), args=([1]))

validation_dataset = tf.data.Dataset.from_generator(raw_data_gen, (tf.float32, tf.uint8), ([None, 1], [None]), args=([2]))

test_dataset = tf.data.Dataset.from_generator(raw_data_gen, (tf.float32, tf.uint8), ([None, 1], [None]), args=([3]))

以这种方式尝试呼叫from_generator()时收到的错误消息是:

TypeError: from_generator() got an unexpected keyword argument 'args'

这里是raw_data_gen函数,尽管我不确定您是否需要此函数,因为我的直觉是问题在于调用from_generator()

def raw_data_gen(train_val_or_test):

    if train_val_or_test == 1:        
        #For every filename collected in the list
        for filename, lab in training_filepath_label_dict.items():
            raw_data, samplerate = soundfile.read(filename)
            try: #assume the audio is stereo, ready to be sliced
                raw_data = raw_data[:,0] #raw_data is a np.array, just take first channel with slice
            except IndexError:
                pass #this must be mono audio
            yield raw_data, lab

    elif train_val_or_test == 2:
        #For every filename collected in the list
        for filename, lab in validation_filepath_label_dict.items():
            raw_data, samplerate = soundfile.read(filename)
            try: #assume the audio is stereo, ready to be sliced
                raw_data = raw_data[:,0] #raw_data is a np.array, just take first channel with slice
            except IndexError:
                pass #this must be mono audio
            yield raw_data, lab

    elif train_val_or_test == 3:
        #For every filename collected in the list
        for filename, lab in test_filepath_label_dict.items():
            raw_data, samplerate = soundfile.read(filename)
            try: #assume the audio is stereo, ready to be sliced
                raw_data = raw_data[:,0] #raw_data is a np.array, just take first channel with slice
            except IndexError:
                pass #this must be mono audio
            yield raw_data, lab

    else:
        print("generator function called with an argument not in [1, 2, 3]")
        raise ValueError()

2 个答案:

答案 0 :(得分:4)

您需要基于raw_data_gen定义一个不带任何参数的新函数。您可以使用lambda关键字来完成此操作。

training_dataset = tf.data.Dataset.from_generator(lambda: raw_data_gen(train_val_or_test=1), (tf.float32, tf.uint8), ([None, 1], [None]))
...

现在,我们正在将一个不带任何参数的函数传递给from_generator,但是它将简单地充当raw_data_gen并将参数设置为1。验证和测试集,分别通过2和3。

答案 1 :(得分:3)

对于 Tensorflow 2.4:

training_dataset = tf.data.Dataset.from_generator(
     raw_data_gen, 
     args=(1), 
     output_types=(tf.float32, tf.uint8), 
     output_shapes=([None, 1], [None]))